qBitDownload-Bot/tubot/torrent/apis.py

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

########################################
# Torrent Server API's implementations #
########################################

# Imports
from http.cookies import SimpleCookie
from aiohttp import ClientSession, ClientResponse, FormData
from tubot.torrent.abc import TorrentAPI
from tubot.torrent.torrents import TorrentFile, TorrentMagnet, TorrentURL
from tubot.torrent.types import ServerTypes, TorrentListBuilder
from tubot.static.functions import validate
from tubot.static.exceptions import AuthError


class qBitTorrent(TorrentAPI):
    """
    qBitTorrent API implementation
    """

    host: str
    username: str
    password: str
    cookie: SimpleCookie | None
    _atype = ServerTypes.qBitTorrent

    def __init__(self, host: str, username: str, password: str) -> None:
        """
        :param host: qBitTorrent remote server adress
        :param username: qBitTorrent remote username
        :param password: qBitTorrent remote password
        """
        super().__init__()
        self.cookie = None
        self.host = host
        self.username = username
        self.password = password

    async def _get(
        self, api: str, cookie: SimpleCookie | None = None
    ) -> ClientResponse:
        """
        Send get request to Torrent server

        :param api: API schema
        :param cookie: Cookies for auth
        """
        async with ClientSession() as session:
            return await session.get(url=f"{self.host}/{api}", cookies=cookie)

    async def _post(
        self,
        api: str,
        cookie: SimpleCookie | None = None,
        data: dict | FormData | None = None,
    ) -> ClientResponse:
        """
        Send post request to Torrent server

        :param api: API schema
        :param cookie: Cookies for auth
        :param data: Request data
        """
        async with ClientSession() as session:
            return await session.post(
                url=f"{self.host}/{api}", cookies=cookie, data=data
            )

    async def auth(self) -> bool:
        """
        Generates cookies for auth
        """
        creds = {"username": self.username, "password": self.password}
        resp = await self._post(api="api/v2/auth/login", data=creds)
        try:
            if resp.status == 200:
                cookies = resp.cookies
                resp = await self._get(api="api/v2/app/version", cookie=cookies)
                if resp.status != 200:
                    raise AuthError("Wrong creds")
                self.cookie = cookies
                return True
        except Exception:
            pass
        return False

    async def upload_file(self, torrent: TorrentFile) -> None:
        await validate(self, "Connection to TorrentServer failed")
        await validate(torrent, "Bad .torrent file")
        bytes = await torrent.getbytes()
        data = FormData()
        data.add_field(
            "torrents",
            bytes,
            filename=torrent.content,
            content_type="application/x-bittorrent",
        )
        data.add_field("savepath", torrent.dest)
        await self._post("api/v2/torrents/add", cookie=self.cookie, data=data)

    async def upload_magnet(self, torrent: TorrentMagnet) -> None:
        await validate(self, "Connection to TorrentServer failed")
        await validate(torrent, "Bad magnet link")
        data = {"urls": torrent.content, "savepath": torrent.dest}
        await self._post("api/v2/torrents/add", cookie=self.cookie, data=data)

    async def upload_url(self, torrent: TorrentURL) -> None:
        await validate(self, "Connection to TorrentServer failed")
        await validate(torrent, "Bad url")
        data = {"urls": torrent.content, "savepath": torrent.dest}
        await self._post("api/v2/torrents/add", cookie=self.cookie, data=data)

    @property
    async def torrent_list(self) -> str:
        await validate(self, "Connection to TorrentServer failed")
        responce = await self._get(
            "api/v2/torrents/info?filter=completed,downloading&sort=progress",
            cookie=self.cookie,
        )
        responce = await responce.json()
        lst = tuple(
            map(lambda i: (i["name"], i["state"], float(i["progress"])), responce)
        )
        lb = TorrentListBuilder()
        for torrent in lst:
            lb.append(torrent)
        return str(lb)

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