qBitDownload-Bot/tubot/dirgetter/getter.py

# -*- coding: utf-8 -*-

####################################
# Directory-Getter implementations #
####################################

# Imports
from tubot.dirgetter.types import GetterTypes
from tubot.dirgetter.abc import DirGetter
from aiofiles.os import listdir
from aiofiles.ospath import isdir
from aiohttp import ClientResponse, ClientSession


class OSGetter(DirGetter):
    """
    Python.os module implementation of DirectoryGetter
    """

    _gtype = GetterTypes.OS
    base_dir: str

    def __init__(self, base_dir: str) -> None:
        """
        :param base_dir: Path to parent directory
        """
        super().__init__()
        self.base_dir = base_dir

    @property
    async def folders(self) -> dict:
        dirs = {}
        ls = await listdir(self.base_dir)
        if len(ls) == 0:
            raise KeyError("No dirs found")
        for item in ls:
            if await isdir(f"{self.base_dir}/{item}"):
                dirs[item] = f"{self.base_dir}/{item}"
        return dirs

    async def __validate__(self) -> bool:
        return await isdir(self.base_dir)


class Jellyfin(DirGetter):
    """
    Jellyfin API implementation of DirectoryGetter
    """

    _gtype = GetterTypes.Jellyfin
    host: str
    token: str

    def __init__(self, host: str, api_token: str) -> None:
        """
        :param host: Adress of Jellyfin server
        :param api_token: Jellyfin API Token for auth
        """
        super().__init__()
        self.host = host
        self.token = api_token

    async def _get(self, api: str) -> ClientResponse:
        async with ClientSession() as session:
            resp = await session.get(f"{self.host}/{api}?api_key={self.token}")
            status = resp.status
            match status:
                case 200:
                    return resp
                case 401:
                    raise ConnectionError("401: Auth error")
                case 403:
                    raise ConnectionError("403: Forbidden")
                case 404:
                    raise ConnectionError("403: Not found")
            raise ConnectionError()

    @property
    async def idx(self) -> str | None:
        resp = await self._get("System/Info")
        json = await resp.json()
        return json["Id"]

    @property
    async def folders(self) -> dict:
        resp = await self._get("Library/VirtualFolders")
        json = await resp.json()
        dirs = {}
        for folder in json:
            dirs[folder["Name"]] = folder["Locations"][0]
        return dirs

    async def __validate__(self) -> bool:
        if await self.idx is not None:
            return True
        return False