diff --git a/media_manager/torrent/download_clients/sabnzbd.py b/media_manager/torrent/download_clients/sabnzbd.py index cc2893b3..d9927336 100644 --- a/media_manager/torrent/download_clients/sabnzbd.py +++ b/media_manager/torrent/download_clients/sabnzbd.py @@ -1,47 +1,23 @@ -import logging - -from media_manager.config import AllEncompassingConfig -from media_manager.indexer.schemas import IndexerQueryResult from media_manager.torrent.download_clients.abstractDownloadClient import ( AbstractDownloadClient, ) +from media_manager.usenet.download_clients.torrentNzbAdapter import TorrentNzbAdapter +from media_manager.usenet.download_clients.sabnzbd import SabnzbdDownloadClient as SabnzbdUsenetDownloadClient +from media_manager.indexer.schemas import IndexerQueryResult from media_manager.torrent.schemas import Torrent, TorrentStatus -import sabnzbd_api - -log = logging.getLogger(__name__) - class SabnzbdDownloadClient(AbstractDownloadClient): - name = "sabnzbd" + """ + Sabnzbd download client abiding by the abstract download client. - DOWNLOADING_STATE = ( - "Downloading", - "Queued", - "Paused", - "Extracting", - "Moving", - "Running", - ) - FINISHED_STATE = ("Completed",) - ERROR_STATE = ("Failed",) - UNKNOWN_STATE = ("Unknown",) + Converts calls to/from compatible calls with usenet.download_clients.sabnzbd.SabnzbdDownloadClient. This allows the + Sabnzbd download client to exist with Usenet/Nzb-first logic, while still only being partially abstracted from the + torrent-first functionality. + """ + name = "sabnzbd" def __init__(self): - self.config = AllEncompassingConfig().torrents.sabnzbd - self.client = sabnzbd_api.SabnzbdClient( - host=self.config.host, - port=str(self.config.port), - api_key=self.config.api_key, - ) - self.client._base_url = f"{self.config.host.rstrip('/')}:{self.config.port}{self.config.base_path}" # the library expects a /sabnzbd prefix for whatever reason - 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 + self.usenet = SabnzbdUsenetDownloadClient() def download_torrent(self, indexer_result: IndexerQueryResult) -> Torrent: """ @@ -50,41 +26,9 @@ def download_torrent(self, indexer_result: IndexerQueryResult) -> Torrent: :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), nzbname=indexer_result.title - ) - if not response["status"]: - error_msg = response - log.error(f"Failed to add NZB to SABnzbd: {error_msg}") - raise RuntimeError(f"Failed to add NZB to SABnzbd: {error_msg}") + nzb = self.usenet.download_nzb(indexer_result) + return TorrentNzbAdapter.nzb_to_torrent(nzb) - # 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: """ @@ -93,13 +37,8 @@ def remove_torrent(self, torrent: Torrent, delete_data: bool = False) -> None: :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 + nzb = TorrentNzbAdapter.torrent_to_nzb(torrent) + self.usenet.remove_nzb(nzb, delete_data) def pause_torrent(self, torrent: Torrent) -> None: """ @@ -107,13 +46,8 @@ def pause_torrent(self, torrent: Torrent) -> None: :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 + nzb = TorrentNzbAdapter.torrent_to_nzb(torrent) + self.usenet.pause_nzb(nzb) def resume_torrent(self, torrent: Torrent) -> None: """ @@ -121,13 +55,8 @@ def resume_torrent(self, torrent: Torrent) -> None: :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 + nzb = TorrentNzbAdapter.torrent_to_nzb(torrent) + return self.usenet.resume_nzb(nzb) def get_torrent_status(self, torrent: Torrent) -> TorrentStatus: """ @@ -136,25 +65,6 @@ def get_torrent_status(self, torrent: Torrent) -> TorrentStatus: :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) - log.debug("SABnzbd response: %s", response) - status = response["queue"]["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 + nzb = TorrentNzbAdapter.torrent_to_nzb(torrent) + status = self.usenet.get_nzb_status(nzb) + return TorrentNzbAdapter.convert_status(status) diff --git a/media_manager/torrent/schemas.py b/media_manager/torrent/schemas.py index dfcca717..9230e7ce 100644 --- a/media_manager/torrent/schemas.py +++ b/media_manager/torrent/schemas.py @@ -25,9 +25,13 @@ class QualityStrings(Enum): class TorrentStatus(Enum): finished = 1 + # Torrent is finished downloading and ready to move. downloading = 2 + # Torrent is downloading. error = 3 + # Torrent failed to download. unknown = 4 + # Unable to obtain status of torrent. class Torrent(BaseModel): diff --git a/media_manager/usenet/download_clients/__init__.py b/media_manager/usenet/download_clients/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/media_manager/usenet/download_clients/abstractDownloadClient.py b/media_manager/usenet/download_clients/abstractDownloadClient.py new file mode 100644 index 00000000..3e069271 --- /dev/null +++ b/media_manager/usenet/download_clients/abstractDownloadClient.py @@ -0,0 +1,64 @@ +from abc import ABC, abstractmethod + +from media_manager.indexer.schemas import IndexerQueryResult +from media_manager.usenet.schemas import NzbStatus, Nzb + + +class AbstractDownloadClient(ABC): + """ + Abstract base class for download clients. + Defines the interface that all download clients must implement. + """ + + @property + @abstractmethod + def name(self) -> str: + pass + + @abstractmethod + def download_nzb(self, nzb: IndexerQueryResult) -> Nzb: + """ + Add a nzb to the download client and return the nzb object. + + :param nzb: The indexer query result of the nzb file to download. + :return: The nzb object with calculated hash and initial status. + """ + pass + + @abstractmethod + def remove_nzb(self, nzb: Nzb, delete_data: bool = False) -> None: + """ + Remove a nzb from the download client. + + :param nzb: The nzb to remove. + :param delete_data: Whether to delete the downloaded data. + """ + pass + + @abstractmethod + def get_nzb_status(self, nzb: Nzb) -> NzbStatus: + """ + Get the status of a specific nzb. + + :param nzb: The nzb to get the status of. + :return: The status of the nzb. + """ + pass + + @abstractmethod + def pause_nzb(self, nzb: Nzb) -> None: + """ + Pause a nzb download. + + :param nzb: The nzb to pause. + """ + pass + + @abstractmethod + def resume_nzb(self, nzb: Nzb) -> None: + """ + Resume a nzb download. + + :param nzb: The nzb to resume. + """ + pass diff --git a/media_manager/usenet/download_clients/sabnzbd.py b/media_manager/usenet/download_clients/sabnzbd.py new file mode 100644 index 00000000..af502d99 --- /dev/null +++ b/media_manager/usenet/download_clients/sabnzbd.py @@ -0,0 +1,175 @@ +import logging + +from media_manager.config import AllEncompassingConfig +from media_manager.indexer.schemas import IndexerQueryResult +from media_manager.usenet.schemas import ( + Nzb, + NzbStatus, + sabnzbd_to_nzb_status +) + +import sabnzbd_api + +log = logging.getLogger(__name__) + + +class SabnzbdDownloadClient: + name = "sabnzbd" + + def __init__(self): + self.config = AllEncompassingConfig().torrents.sabnzbd + self.client = sabnzbd_api.SabnzbdClient( + host=self.config.host, + port=str(self.config.port), + api_key=self.config.api_key, + ) + self.client._base_url = f"{self.config.host.rstrip('/')}:{self.config.port}{self.config.base_path}" # the library expects a /sabnzbd prefix for whatever reason + 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_nzb(self, indexer_result: IndexerQueryResult) -> Nzb: + """ + Add a NZB to SABnzbd and return the Nzb 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), nzbname=indexer_result.title + ) + if not response["status"]: + error_msg = response + 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 + nzb = Nzb( + status=NzbStatus.unknown, + title=indexer_result.title, + quality=indexer_result.quality, + imported=False, + hash=nzo_id, + ) + + # Get initial status from SABnzbd + nzb.status = self.get_nzb_status(nzb) + + return nzb + + except Exception as e: + log.error(f"Failed to download NZB {indexer_result.title}: {e}") + raise + + def remove_nzb(self, nzb: Nzb, delete_data: bool = False) -> None: + """ + Remove a nzb from SABnzbd. + + :param nzb: The nzb to remove. + :param delete_data: Whether to delete the downloaded files. + """ + log.info(f"Removing nzb: {nzb.title} (Delete data: {delete_data})") + try: + self.client.delete_job(nzo_id=nzb.hash, delete_files=delete_data) + log.info(f"Successfully removed nzb: {nzb.title}") + except Exception as e: + log.error(f"Failed to remove nzb {nzb.title}: {e}") + raise + + def pause_nzb(self, nzb: Nzb) -> None: + """ + Pause a nzb in SABnzbd. + + :param nzb: The nzb to pause. + """ + log.info(f"Pausing nzb: {nzb.title}") + try: + self.client.pause_job(nzo_id=nzb.hash) + log.info(f"Successfully paused nzb: {nzb.title}") + except Exception as e: + log.error(f"Failed to pause nzb {nzb.title}: {e}") + raise + + def resume_nzb(self, nzb: Nzb) -> None: + """ + Resume a paused nzb in SABnzbd. + + :param nzb: The nzb to resume. + """ + log.info(f"Resuming nzb: {nzb.title}") + try: + self.client.resume_job(nzo_id=nzb.hash) + log.info(f"Successfully resumed nzb: {nzb.title}") + except Exception as e: + log.error(f"Failed to resume nzb {nzb.title}: {e}") + raise + + + def get_nzb_status(self, nzb: Nzb) -> NzbStatus: + """ + Get the status of a specific download from SABnzbd. + + :param nzb: The nzb to get the status of. + :return: The status of the nzb. + """ + # Check the queue for in progress downloads + status = self._get_in_progress_status(nzb) + if status is not NzbStatus.unknown: + return status + # Check the history for processing or completed downloads + status = self._get_finished_status(nzb) + if status is not NzbStatus.unknown: + return status + + log.warning(f"Could not find any downloads for nzb: {nzb.title}") + return status + + def _get_in_progress_status(self, nzb: Nzb) -> NzbStatus: + """ + Check SABnzbd in progress downloads for torrent status. + + :param torrent: The torrent for which to get the status. + """ + log.info(f"Checking in progress downloads for status: {nzb.title}") + response = self.client.get_downloads(nzo_ids=nzb.hash) + log.debug("SABnzbd queue response: %s", response) + nzb_slots = response["queue"]["slots"] + if len(nzb_slots) >= 1: + status = nzb_slots[0]["status"] + log.info(f"Status for in progress nzb {nzb.title}: {status}") + return sabnzbd_to_nzb_status(status) + + log.info(f"Didn't find any downloads in progress for nzb: {nzb.title}") + return NzbStatus.unknown + + def _get_finished_status(self, nzb: Nzb) -> NzbStatus: + """ + Check SABnzbd finished downloads for nzb status. + + :param nzb: The nzb for which to get the status. + """ + log.info(f"Checking completed downloads for status: {nzb.title}") + response = self.client.get_history(nzo_ids=nzb.hash) + log.debug("SABnzbd history response: %s", response) + nzb_slots = response["history"]["slots"] + if len(nzb_slots) >= 1: + status = nzb_slots[0]["status"] + log.info(f"Status for downloaded nzb {nzb.title}: {status}") + return sabnzbd_to_nzb_status(status) + + log.info(f"Didn't find any completed downloads for nzb: {nzb.title}") + return NzbStatus.unknown diff --git a/media_manager/usenet/download_clients/torrentNzbAdapter.py b/media_manager/usenet/download_clients/torrentNzbAdapter.py new file mode 100644 index 00000000..77d31ce9 --- /dev/null +++ b/media_manager/usenet/download_clients/torrentNzbAdapter.py @@ -0,0 +1,74 @@ +""" +Adapts Torrent objects to Nzb objects, and vice versa. + +Allows for Usenet functions to be partially abstracted from Torrent functions, +pending a full breakout. +""" +import logging +from media_manager.usenet.schemas import Nzb, NzbStatus +from media_manager.torrent.schemas import Torrent, TorrentStatus + +log = logging.getLogger(__name__) + +class TorrentNzbAdapter: + name = "TorrentNzbAdapter" + + @staticmethod + def nzb_to_torrent(nzb: Nzb) -> Torrent: + """ + Convert a Nzb object to a Torrent object. + + :param nzb: The Nzb object to convert. + :return: The converted Torrent object. + """ + log.debug(f"Converting Nzb to Torrent: {nzb}") + status = TorrentNzbAdapter.convert_status(nzb.status) if nzb.status else TorrentStatus.unknown + torrent = Torrent( + status=status, + title=nzb.title, + quality=nzb.quality, + imported=nzb.imported, + hash=nzb.hash, + usenet=True, + ) + + return torrent + + @staticmethod + def torrent_to_nzb(torrent: Torrent) -> Nzb: + """ + Convert a Torrent object to Nzb object. + + :param torrent: The Torrent object to convert. + :return: The converted Nzb object. + """ + log.debug(f"Converting Torrent to Nzb: {torrent}") + status = TorrentNzbAdapter.convert_status(torrent.status) if torrent.status else NzbStatus.unknown + nzb = Nzb( + status=status, + title=torrent.title, + quality=torrent.quality, + imported=torrent.imported, + hash=torrent.hash, + ) + + return nzb + + @staticmethod + def convert_status(status_obj: NzbStatus or TorrentStatus) -> NzbStatus or TorrentStatus: + """ + Converts NzbStatus to TorrentStatus or vice versa. NzbStatus 5 (processing) is folded + into TorrentStatus 2 (downloading). + + :param status_obj: The NzbStatus or TorrentStatus object to convert. + :return: The converted NzbStatus or TorrentStatus object. + """ + log.debug(f"Converting status object {status_obj} of type {type(status_obj)}") + log.debug(f"status object value is {status_obj.value}") + if isinstance(status_obj, NzbStatus) and status_obj.value == 5: + return TorrentStatus(2) + elif isinstance(status_obj, NzbStatus): + return TorrentStatus(status_obj.value) + elif isinstance(status_obj, TorrentStatus) : + return NzbStatus(status_obj.value) + diff --git a/media_manager/usenet/schemas.py b/media_manager/usenet/schemas.py new file mode 100644 index 00000000..6b558a27 --- /dev/null +++ b/media_manager/usenet/schemas.py @@ -0,0 +1,116 @@ +import typing +import uuid +from enum import Enum + +from pydantic import ConfigDict, BaseModel, Field + +from media_manager.torrent.schemas import Quality + +NzbId = typing.NewType("NzbId", uuid.UUID) + + +class NzbStatus(Enum): + """ + Defines MediaManager's simplified Nzb statuses. + """ + finished = 1 + # File(s) are finished downloading and processing, and are ready to move. + downloading = 2 + # File(s) are downloading. + error = 3 + # File(s) failed to download. + unknown = 4 + # Unable to obtain status of download. + processing = 5 + # File(s) are downloaded and processing before completion. + + +class Nzb(BaseModel): + model_config = ConfigDict(from_attributes=True) + + id: NzbId = Field(default_factory=uuid.uuid4) + status: NzbStatus + title: str + quality: Quality + imported: bool + hash: str + usenet: bool = True + + +class SabnzbdStatus(Enum): + """ + Defines all Sabnzbd nzb statuses. + + https://sabnzbd.org/wiki/extra/queue-history-searching + """ + ########################### + # in progress nzb download statuses. + ########################## + Grabbing = "Grabbing" + # Getting an NZB from an external site + Downloading = "Downloading" + # Downloading the file + Paused = "Paused" + # Individual download and/or entire download queue is paused + Propagating = "Propagating" + # Download is delayed + ########################### + # downloaded nzb statuses. + ########################## + Completed = "Completed" + # Download and processing is complete + Failed = "Failed" + # Download failed + QuickCheck = "QuickCheck" + # Fast integrity check of download + Verifying = "Verifying" + # Running SFV-based integrity check (by par2) + Repairing = "Repairing" + # Job is being repaired (by par2) + Extracting = "Extracting" + # Extracting the completed download archives + Moving = "Moving" + # Moving the final media file to its destination + Running = "Running" + # Post-processing script is running + ########################### + # statuses for both in progress and downloaded nzbs. + ########################## + Fetching = "Fetching" + # Job is downloading extra par2 files + Queued = "Queued" + # Download is queued, or if finished, waiting for post-processing + Unknown = "Unknown" + # Download status was not found. + + +def sabnzbd_to_nzb_status(status: str or SabnzbdStatus) -> NzbStatus: + """ + Reduces any Sabnzbd status into a simplified Nzb status. + """ + if isinstance(status, str): + status = SabnzbdStatus(status) + + if status in [ + SabnzbdStatus.Downloading, + SabnzbdStatus.Fetching, + SabnzbdStatus.Queued, + SabnzbdStatus.Paused, + SabnzbdStatus.Propagating, + SabnzbdStatus.Grabbing, + ]: + return NzbStatus.downloading + elif status in [ + SabnzbdStatus.QuickCheck, + SabnzbdStatus.Verifying, + SabnzbdStatus.Extracting, + SabnzbdStatus.Moving, + SabnzbdStatus.Running, + ]: + return NzbStatus.processing + elif status is SabnzbdStatus.Completed: + return NzbStatus.finished + elif status is SabnzbdStatus.Failed: + return NzbStatus.error + else: + return NzbStatus.unknown