From e310e5cc9c816e0aa197433f6033a8abc09f122d Mon Sep 17 00:00:00 2001 From: fish-portable Date: Wed, 25 Oct 2023 21:45:31 +0800 Subject: [PATCH] feat(transmission): add transmission downloader and model --- .../module/downloader/client/tr_downloader.py | 173 +++++++++++++++++- .../src/module/downloader/download_client.py | 4 + backend/src/module/models/config.py | 19 ++ 3 files changed, 191 insertions(+), 5 deletions(-) diff --git a/backend/src/module/downloader/client/tr_downloader.py b/backend/src/module/downloader/client/tr_downloader.py index 29591525..0219bf64 100644 --- a/backend/src/module/downloader/client/tr_downloader.py +++ b/backend/src/module/downloader/client/tr_downloader.py @@ -1,13 +1,23 @@ +import logging import httpx +import base64 + +logger = logging.getLogger(__name__) class TrDownloader: def __init__(self, host, username, password, ssl): - self.host = host + self.host = host if "://" in host else "http://" + host self.username = username self.password = password self.ssl = ssl - self._client = None + self.authkey = base64.b64encode(f"{self.username}:{self.password}".encode()).decode() + + self._client = httpx.AsyncClient( + base_url=self.host, + auth=(self.username, self.password), + timeout=5, + ) async def __aenter__(self): self._client = httpx.AsyncClient( @@ -18,9 +28,162 @@ class TrDownloader: return self async def __aexit__(self, exc_type, exc_val, exc_tb): + self.logout() await self._client.aclose() async def auth(self): - resp = await self._client.get("/transmission/rpc") - resp.raise_for_status() - return resp + # NOTE: Transmission will return 409 when first login + if self.username and self.password: + self._client.headers.update({"Authorization": f"Basic {self.authkey}"}) + + resp = await self._client.post("/transmission/rpc") + + if resp.status_code == 409 and 'X-Transmission-Session-Id' in resp.headers: + self._client.headers.update({"X-Transmission-Session-Id": resp.headers['X-Transmission-Session-Id']}) + resp = await self._client.post("/transmission/rpc") + elif resp.status_code == 401: + logger.error("Transmission: Authentication failed") + return False + + return resp.status_code == 200 + + def logout(self): + self._client.headers.pop("Authorization") + + async def check_host(self): + try: + resp = await self._client.get("/transmission/web/") + return resp.status_code == 200 + except httpx.RequestError: + return False + + async def add_torrent(self, download_link=None, torrent_path=None, save_path=None, **kwargs): + + if not download_link and not torrent_path: + # WARNING: Regard no torrent as success + return True + request_data = { + "method": "torrent-add", + "arguments": { + "download-dir": save_path, + "paused": False, + **kwargs + } + } + + if torrent_path: + try: + with open(torrent_path, 'rb') as file: + file_content = file.read() + metainfo = base64.b64encode(file_content).decode() + except FileNotFoundError: + logger.error(f"File not found: {torrent_path}") + return False + + request_data["arguments"].update({"metainfo": metainfo}) + else: + request_data["arguments"].update({"filename": download_link}) + + resp = await self._client.post("/transmission/rpc", json=request_data) + + return resp.status_code == 200 + + async def add(self, torrent_urls, torrent_files, save_path, category): + result = True + for torrent_url in torrent_urls: + result = result and \ + await self.add_torrent(download_link=torrent_url, save_path=save_path, labels=[category]) + + for torrent_file in torrent_files: + result = result and \ + await self.add_torrent(torrent_path=torrent_file, save_path=save_path) + + return result + + async def delete(self, _hash): + request_data = { + "method": "torrent-remove", + "arguments": { + "ids": [_hash], + "delete-local-data": True + } + } + resp = await self._client.post("/transmission/rpc", json=request_data) + return resp.status_code == 200 + + async def move(self, hashes, new_location): + request_data = { + "method": "torrent-set-location", + "arguments": { + "ids": hashes, + "location": new_location + } + } + resp = await self._client.post("/transmission/rpc", json=request_data) + return resp.status_code == 200 + + async def rename(self, torrent_hash, old_path, new_path) -> bool: + request_data = { + "method": "torrent-rename-path", + "arguments": { + "ids": [torrent_hash], + "path": old_path, + "name": new_path + } + } + resp = await self._client.post("/transmission/rpc", json=request_data) + return resp.status_code == 200 + + async def torrents_info(self, status_filter, category, tag=None): + KEY_MAP = {"hashString": "hash", "downloadDir": "save_path"} + # Map transmission key to qbittorrent + + request_data = { + "method": "torrent-get", + "arguments": { + "fields": ["id", "name", "hashString", "downloadDir", "status", "labels"], + }, + "format": "object" + } + resp = await self._client.post("/transmission/rpc", json=request_data) + data = resp.json() + torrents_info = data['arguments'].get("torrents") + for torrent_info in torrents_info: + for old_key, new_key in KEY_MAP.items(): + torrent_info[new_key] = torrent_info.pop(old_key) + + torrents_info = self._filter_status(torrents_info, status_filter) + if category: + torrents_info = [torrent for torrent in torrents_info if category in torrent["labels"]] + # NOTE: To compatible with qbittorrent api we use category as label + + return torrents_info + + async def set_category(self, torrent_hashes, category): + request_data = { + "method": "torrent-set", + "arguments": { + "ids": torrent_hashes, + "labels": [category] + } + } + + # NOTE: To compatible with qbittorrent api we use category as label + resp = await self._client.post("/transmission/rpc", json=request_data) + return resp.status_code == 200 + + def _filter_status(self, torrents_info, status_filter: str): + """ + Filter torrents by status + Docs: https://github.com/transmission/transmission/blob/main/docs/rpc-spec.md#33-torrent-accessor-torrent-get + """ + if status_filter == "completed": + # We regard torrents queue to seed as completed + return [torrent for torrent in torrents_info if torrent['status'] >= 5] + elif status_filter == "downloading": + return [torrent for torrent in torrents_info if torrent['status'] == 4] + elif status_filter == "inactive": + return [torrent for torrent in torrents_info if torrent['status'] <= 3] + + return torrents_info + diff --git a/backend/src/module/downloader/download_client.py b/backend/src/module/downloader/download_client.py index 830d9991..88d39f78 100644 --- a/backend/src/module/downloader/download_client.py +++ b/backend/src/module/downloader/download_client.py @@ -14,6 +14,10 @@ def getClient(): if settings.downloader.type == "qbittorrent": from .client.qb_downloader import QbDownloader return QbDownloader + elif type == "transmission": + from .client.tr_downloader import TrDownloader + + return TrDownloader else: logger.error(f"[Downloader] Unsupported downloader type: {type}") raise Exception(f"Unsupported downloader type: {type}") diff --git a/backend/src/module/models/config.py b/backend/src/module/models/config.py index 49044fa3..34a4385f 100644 --- a/backend/src/module/models/config.py +++ b/backend/src/module/models/config.py @@ -11,6 +11,14 @@ class Program(BaseModel): class Downloader(BaseModel): + host_: str + username_: str + password_: str + path: str + ssl: bool + + +class QbDownloader(Downloader): type: str = Field("qbittorrent", description="Downloader type") host_: str = Field("172.17.0.1:8080", alias="host", description="Downloader host") username_: str = Field("admin", alias="username", description="Downloader username") @@ -33,6 +41,17 @@ class Downloader(BaseModel): return expandvars(self.password_) +class TrDownloader(Downloader): + type: str = Field("transmission", description="Downloader type") + host_: str = Field("172.17.0.1:9091", alias="host", description="Downloader host") + username_: str = Field("admin", alias="username", description="Downloader username") + password_: str = Field( + "admin", alias="password", description="Downloader password" + ) + path: str = Field("/downloads/Bangumi", description="Downloader path") + ssl: bool = Field(False, description="Downloader ssl") + + class RSSParser(BaseModel): enable: bool = Field(True, description="Enable RSS parser") filter: list[str] = Field(["720", r"\d+-\d"], description="Filter")