From 28349641c02a81dabf5b8dbdf9d6cd0cecd93917 Mon Sep 17 00:00:00 2001 From: maxDorninger <97409287+maxDorninger@users.noreply.github.com> Date: Wed, 9 Jul 2025 23:54:51 +0200 Subject: [PATCH] format files --- .../abstractDownloadClient.py | 22 ++- .../torrent/download_clients/qbittorrent.py | 22 ++- .../torrent/download_clients/sabnzbd.py | 163 ++++++++++++++++++ media_manager/torrent/service.py | 9 +- pyproject.toml | 1 + uv.lock | 15 ++ 6 files changed, 221 insertions(+), 11 deletions(-) create mode 100644 media_manager/torrent/download_clients/sabnzbd.py diff --git a/media_manager/torrent/download_clients/abstractDownloadClient.py b/media_manager/torrent/download_clients/abstractDownloadClient.py index 2a31282..2a8e20f 100644 --- a/media_manager/torrent/download_clients/abstractDownloadClient.py +++ b/media_manager/torrent/download_clients/abstractDownloadClient.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from media_manager.indexer.schemas import IndexerQueryResult -from media_manager.torrent.schemas import TorrentId, TorrentStatus, Torrent +from media_manager.torrent.schemas import TorrentStatus, Torrent class AbstractDownloadClient(ABC): @@ -38,4 +38,22 @@ class AbstractDownloadClient(ABC): :param torrent: The torrent to get the status of. :return: The status of the torrent. """ - pass \ No newline at end of file + pass + + @abstractmethod + def pause_torrent(self, torrent: Torrent) -> None: + """ + Pause a torrent download. + + :param torrent: The torrent to pause. + """ + pass + + @abstractmethod + def resume_torrent(self, torrent: Torrent) -> None: + """ + Resume a torrent download. + + :param torrent: The torrent to resume. + """ + pass diff --git a/media_manager/torrent/download_clients/qbittorrent.py b/media_manager/torrent/download_clients/qbittorrent.py index fdea9cd..86090b7 100644 --- a/media_manager/torrent/download_clients/qbittorrent.py +++ b/media_manager/torrent/download_clients/qbittorrent.py @@ -8,8 +8,10 @@ from pydantic_settings import BaseSettings, SettingsConfigDict from media_manager.config import BasicConfig from media_manager.indexer.schemas import IndexerQueryResult -from media_manager.torrent.download_clients.abstractDownloadClient import AbstractDownloadClient -from media_manager.torrent.schemas import TorrentId, TorrentStatus, Torrent +from media_manager.torrent.download_clients.abstractDownloadClient import ( + AbstractDownloadClient, +) +from media_manager.torrent.schemas import TorrentStatus, Torrent log = logging.getLogger(__name__) @@ -66,7 +68,9 @@ class QbittorrentDownloadClient(AbstractDownloadClient): """ log.info(f"Attempting to download torrent: {indexer_result.title}") - torrent_filepath = BasicConfig().torrent_directory / f"{indexer_result.title}.torrent" + torrent_filepath = ( + BasicConfig().torrent_directory / f"{indexer_result.title}.torrent" + ) if torrent_filepath.exists(): log.warning(f"Torrent already exists: {torrent_filepath}") @@ -74,7 +78,9 @@ class QbittorrentDownloadClient(AbstractDownloadClient): with open(torrent_filepath, "rb") as file: content = file.read() decoded_content = bencoder.decode(content) - torrent_hash = hashlib.sha1(bencoder.encode(decoded_content[b"info"])).hexdigest() + torrent_hash = hashlib.sha1( + bencoder.encode(decoded_content[b"info"]) + ).hexdigest() else: # Download the torrent file with open(torrent_filepath, "wb") as file: @@ -97,7 +103,9 @@ class QbittorrentDownloadClient(AbstractDownloadClient): try: self.api_client.auth_log_in() answer = self.api_client.torrents_add( - category="MediaManager", torrent_files=content, save_path=indexer_result.title + category="MediaManager", + torrent_files=content, + save_path=indexer_result.title, ) finally: self.api_client.auth_log_out() @@ -134,7 +142,9 @@ class QbittorrentDownloadClient(AbstractDownloadClient): log.info(f"Removing torrent: {torrent.title}") try: self.api_client.auth_log_in() - self.api_client.torrents_delete(torrent_hashes=torrent.hash, delete_files=delete_data) + self.api_client.torrents_delete( + torrent_hashes=torrent.hash, delete_files=delete_data + ) finally: self.api_client.auth_log_out() diff --git a/media_manager/torrent/download_clients/sabnzbd.py b/media_manager/torrent/download_clients/sabnzbd.py new file mode 100644 index 0000000..0fbfb48 --- /dev/null +++ b/media_manager/torrent/download_clients/sabnzbd.py @@ -0,0 +1,163 @@ +import logging +from pydantic_settings import BaseSettings, SettingsConfigDict + +from media_manager.indexer.schemas import IndexerQueryResult +from media_manager.torrent.download_clients.abstractDownloadClient import ( + AbstractDownloadClient, +) +from media_manager.torrent.schemas import Torrent, TorrentStatus +import sabnzbd_api + +log = logging.getLogger(__name__) + + +class SabnzbdConfig(BaseSettings): + model_config = SettingsConfigDict(env_prefix="SABNZBD_") + host: str = "localhost" + port: int = 8080 + api_key: str = "" + + +class SabnzbdDownloadClient(AbstractDownloadClient): + DOWNLOADING_STATE = ( + "Downloading", + "Queued", + "Paused", + "Extracting", + "Moving", + "Running", + ) + FINISHED_STATE = ("Completed",) + ERROR_STATE = ("Failed",) + UNKNOWN_STATE = ("Unknown",) + + def __init__(self): + self.config = SabnzbdConfig() + self.client = sabnzbd_api.SabnzbdClient( + host=self.config.host, + port=str(self.config.port), + api_key=self.config.api_key, + ) + try: + # Test connection + version = self.client.version() + + log.info(f"Successfully connected to SABnzbd version: {version}") + except Exception as e: + log.error(f"Failed to connect to SABnzbd: {e}") + raise + + def download_torrent(self, indexer_result: IndexerQueryResult) -> Torrent: + """ + Add a NZB/torrent to SABnzbd and return the torrent object. + + :param indexer_result: The indexer query result of the NZB file to download. + :return: The torrent object with calculated hash and initial status. + """ + log.info(f"Attempting to download NZB: {indexer_result.title}") + + try: + # Add NZB to SABnzbd queue + response = self.client.add_uri( + url=str(indexer_result.download_url), + ) + if not response["status"] == "True": + error_msg = response.get("error", "Unknown error") + log.error(f"Failed to add NZB to SABnzbd: {error_msg}") + raise RuntimeError(f"Failed to add NZB to SABnzbd: {error_msg}") + + # Generate a hash for the NZB (using title and download URL) + nzo_id = response["nzo_ids"][0] + + log.info(f"Successfully added NZB: {indexer_result.title}") + + # Create and return torrent object + torrent = Torrent( + status=TorrentStatus.unknown, + title=indexer_result.title, + quality=indexer_result.quality, + imported=False, + hash=nzo_id, + usenet=True, + ) + + # Get initial status from SABnzbd + torrent.status = self.get_torrent_status(torrent) + + return torrent + + except Exception as e: + log.error(f"Failed to download NZB {indexer_result.title}: {e}") + raise + + def remove_torrent(self, torrent: Torrent, delete_data: bool = False) -> None: + """ + Remove a torrent from SABnzbd. + + :param torrent: The torrent to remove. + :param delete_data: Whether to delete the downloaded files. + """ + log.info(f"Removing torrent: {torrent.title} (Delete data: {delete_data})") + try: + self.client.delete_job(nzo_id=torrent.hash, delete_files=delete_data) + log.info(f"Successfully removed torrent: {torrent.title}") + except Exception as e: + log.error(f"Failed to remove torrent {torrent.title}: {e}") + raise + + def pause_torrent(self, torrent: Torrent) -> None: + """ + Pause a torrent in SABnzbd. + + :param torrent: The torrent to pause. + """ + log.info(f"Pausing torrent: {torrent.title}") + try: + self.client.pause_job(nzo_id=torrent.hash) + log.info(f"Successfully paused torrent: {torrent.title}") + except Exception as e: + log.error(f"Failed to pause torrent {torrent.title}: {e}") + raise + + def resume_torrent(self, torrent: Torrent) -> None: + """ + Resume a paused torrent in SABnzbd. + + :param torrent: The torrent to resume. + """ + log.info(f"Resuming torrent: {torrent.title}") + try: + self.client.resume_job(nzo_id=torrent.hash) + log.info(f"Successfully resumed torrent: {torrent.title}") + except Exception as e: + log.error(f"Failed to resume torrent {torrent.title}: {e}") + raise + + def get_torrent_status(self, torrent: Torrent) -> TorrentStatus: + """ + Get the status of a specific download from SABnzbd. + + :param torrent: The torrent to get the status of. + :return: The status of the torrent. + """ + log.info(f"Fetching status for download: {torrent.title}") + response = self.client.get_downloads(nzo_ids=torrent.hash) + status = response["queue"][0]["status"] + log.info(f"Download status for NZB {torrent.title}: {status}") + return self._map_status(status) + + def _map_status(self, sabnzbd_status: str) -> TorrentStatus: + """ + Map SABnzbd status to TorrentStatus. + + :param sabnzbd_status: The status from SABnzbd. + :return: The corresponding TorrentStatus. + """ + if sabnzbd_status in self.DOWNLOADING_STATE: + return TorrentStatus.downloading + elif sabnzbd_status in self.FINISHED_STATE: + return TorrentStatus.finished + elif sabnzbd_status in self.ERROR_STATE: + return TorrentStatus.error + else: + return TorrentStatus.unknown diff --git a/media_manager/torrent/service.py b/media_manager/torrent/service.py index 83dcd38..f6eddaa 100644 --- a/media_manager/torrent/service.py +++ b/media_manager/torrent/service.py @@ -1,10 +1,9 @@ import logging -from media_manager.config import BasicConfig from media_manager.indexer.schemas import IndexerQueryResult from media_manager.torrent.download_clients.qbittorrent import QbittorrentDownloadClient from media_manager.torrent.repository import TorrentRepository -from media_manager.torrent.schemas import Torrent, TorrentStatus, TorrentId +from media_manager.torrent.schemas import Torrent, TorrentId from media_manager.tv.schemas import SeasonFile, Show from media_manager.movies.schemas import Movie @@ -12,7 +11,11 @@ log = logging.getLogger(__name__) class TorrentService: - def __init__(self, torrent_repository: TorrentRepository, download_client: QbittorrentDownloadClient = None): + def __init__( + self, + torrent_repository: TorrentRepository, + download_client: QbittorrentDownloadClient = None, + ): self.torrent_repository = torrent_repository self.download_client = download_client or QbittorrentDownloadClient() diff --git a/pyproject.toml b/pyproject.toml index aabf1f0..f52b296 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,4 +31,5 @@ dependencies = [ "pytest>=8.4.0", "pillow>=11.2.1", "pillow-avif-plugin>=1.5.2", + "sabnzbd-api>=0.1.2", ] diff --git a/uv.lock b/uv.lock index 57b2d52..20a1286 100644 --- a/uv.lock +++ b/uv.lock @@ -642,6 +642,7 @@ dependencies = [ { name = "python-json-logger" }, { name = "qbittorrent-api" }, { name = "requests" }, + { name = "sabnzbd-api" }, { name = "sqlalchemy" }, { name = "starlette" }, { name = "tmdbsimple" }, @@ -673,6 +674,7 @@ requires-dist = [ { name = "python-json-logger", specifier = ">=3.3.0" }, { name = "qbittorrent-api", specifier = ">=2025.5.0" }, { name = "requests", specifier = ">=2.32.3" }, + { name = "sabnzbd-api", specifier = ">=0.1.2" }, { name = "sqlalchemy", specifier = ">=2.0.41" }, { name = "starlette", specifier = ">=0.46.2" }, { name = "tmdbsimple", specifier = ">=2.9.1" }, @@ -1092,6 +1094,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b6/97/5a4b59697111c89477d20ba8a44df9ca16b41e737fa569d5ae8bff99e650/rpds_py-0.25.1-cp313-cp313t-win_amd64.whl", hash = "sha256:401ca1c4a20cc0510d3435d89c069fe0a9ae2ee6495135ac46bdd49ec0495763", size = 232218, upload-time = "2025-05-21T12:44:40.512Z" }, ] +[[package]] +name = "sabnzbd-api" +version = "0.1.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "httpx" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a5/ab/a0cf6a4bc977afd60a7f9846c1805cf709db73feef7705ee0c4397924f48/sabnzbd-api-0.1.2.tar.gz", hash = "sha256:1bb0defcb1aa19333f717a63464fbdc8b8a748a00ca4289be5c7496f045d339f", size = 7529, upload-time = "2025-02-20T19:36:03.397Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/14/06/b9a7135a8fe164fe928265647694c0963a6e6a5439228c915ae866d13586/sabnzbd_api-0.1.2-py3-none-any.whl", hash = "sha256:65ab9aecda300e574c5074b0ae4b971f61346d89d9e3f80393dc72c712ac3797", size = 8300, upload-time = "2025-02-20T19:35:58.948Z" }, +] + [[package]] name = "shellingham" version = "1.5.4"