import asyncio import logging from module.conf import settings from module.models import Bangumi, Torrent from module.network import RequestContent from .path import TorrentPath logger = logging.getLogger(__name__) class DownloadClient(TorrentPath): """Unified async download client. Wraps qBittorrent, Aria2, or MockDownloader behind a common interface. Intended to be used as an async context manager; authentication is performed on ``__aenter__`` and the session is closed on ``__aexit__``. """ def __init__(self): super().__init__() self.client = self.__getClient() self.authed = False @staticmethod def __getClient(): """Instantiate the configured downloader client (qbittorrent | aria2 | mock).""" downloader_type = settings.downloader.type host = settings.downloader.host username = settings.downloader.username password = settings.downloader.password ssl = settings.downloader.ssl if downloader_type == "qbittorrent": from .client.qb_downloader import QbDownloader return QbDownloader(host, username, password, ssl) elif downloader_type == "aria2": from .client.aria2_downloader import Aria2Downloader return Aria2Downloader(host, username, password) elif downloader_type == "mock": from .client.mock_downloader import MockDownloader logger.debug("[Downloader] Using MockDownloader for local development") return MockDownloader() else: logger.error("[Downloader] Unsupported downloader type: %s", downloader_type) raise Exception(f"Unsupported downloader type: {downloader_type}") async def __aenter__(self): if not self.authed: await self.auth() if not self.authed: raise ConnectionError("Download client authentication failed") else: logger.error("[Downloader] Already authed.") return self async def __aexit__(self, exc_type, exc_val, exc_tb): if self.authed: await self.client.logout() self.authed = False async def auth(self): self.authed = await self.client.auth() if self.authed: logger.debug("[Downloader] Authed.") else: logger.error("[Downloader] Auth failed.") async def check_host(self): return await self.client.check_host() async def init_downloader(self): """Apply required qBittorrent RSS preferences and create the Bangumi category.""" prefs = { "rss_auto_downloading_enabled": True, "rss_max_articles_per_feed": 500, "rss_processing_enabled": True, "rss_refresh_interval": 30, } await self.client.prefs_init(prefs=prefs) # Category creation may fail if it already exists (HTTP 409) or network issues try: await self.client.add_category("BangumiCollection") except Exception as e: logger.debug( "[Downloader] Could not add category (may already exist): %s", e ) if settings.downloader.path == "": prefs = await self.client.get_app_prefs() settings.downloader.path = self._join_path(prefs["save_path"], "Bangumi") async def set_rule(self, data: Bangumi): """Create or update a qBittorrent RSS auto-download rule for one bangumi entry.""" data.rule_name = self._rule_name(data) data.save_path = self._gen_save_path(data) rule = { "enable": True, "mustContain": data.title_raw, "mustNotContain": "|".join(data.filter), "useRegex": True, "episodeFilter": "", "smartFilter": False, "previouslyMatchedEpisodes": [], "affectedFeeds": data.rss_link, "ignoreDays": 0, "lastMatch": "", "addPaused": False, "assignedCategory": "Bangumi", "savePath": data.save_path, } await self.client.rss_set_rule(rule_name=data.rule_name, rule_def=rule) data.added = True logger.info( f"[Downloader] Add {data.official_title} Season {data.season} to auto download rules." ) async def set_rules(self, bangumi_info: list[Bangumi]): logger.debug("[Downloader] Start adding rules.") await asyncio.gather(*[self.set_rule(info) for info in bangumi_info]) logger.debug("[Downloader] Finished.") async def get_torrent_info( self, category="Bangumi", status_filter="completed", tag=None ): return await self.client.torrents_info( status_filter=status_filter, category=category, tag=tag ) async def get_torrent_files(self, torrent_hash: str): return await self.client.torrents_files(torrent_hash=torrent_hash) async def rename_torrent_file( self, _hash, old_path, new_path, verify: bool = True ) -> bool: result = await self.client.torrents_rename_file( torrent_hash=_hash, old_path=old_path, new_path=new_path, verify=verify ) if result: logger.info(f"{old_path} >> {new_path}") else: logger.debug("[Downloader] Rename failed: %s >> %s", old_path, new_path) return result async def delete_torrent(self, hashes, delete_files: bool = True): await self.client.torrents_delete(hashes, delete_files=delete_files) logger.info("[Downloader] Remove torrents.") async def pause_torrent(self, hashes: str): await self.client.torrents_pause(hashes) async def resume_torrent(self, hashes: str): await self.client.torrents_resume(hashes) async def add_torrent(self, torrent: Torrent | list, bangumi: Bangumi) -> bool: """Download a torrent (or list of torrents) for the given bangumi entry. Handles both magnet links and .torrent file URLs, fetching file bytes when necessary. Tags each torrent with ``ab:`` for later episode-offset lookup during rename. """ if not bangumi.save_path: bangumi.save_path = self._gen_save_path(bangumi) async with RequestContent() as req: if isinstance(torrent, list): if len(torrent) == 0: logger.debug( "[Downloader] No torrent found: %s", bangumi.official_title ) return False if "magnet" in torrent[0].url: torrent_url = [t.url for t in torrent] torrent_file = None else: torrent_file = await asyncio.gather( *[req.get_content(t.url) for t in torrent] ) # Filter out None values (failed fetches) torrent_file = [f for f in torrent_file if f is not None] if not torrent_file: logger.warning( f"[Downloader] Failed to fetch torrent files for: {bangumi.official_title}" ) return False torrent_url = None else: if "magnet" in torrent.url: torrent_url = torrent.url torrent_file = None else: torrent_file = await req.get_content(torrent.url) if torrent_file is None: logger.warning( f"[Downloader] Failed to fetch torrent file for: {bangumi.official_title}" ) return False torrent_url = None # Create tag with bangumi_id for offset lookup during rename tags = f"ab:{bangumi.id}" if bangumi.id else None try: if await self.client.add_torrents( torrent_urls=torrent_url, torrent_files=torrent_file, save_path=bangumi.save_path, category="Bangumi", tags=tags, ): logger.debug("[Downloader] Add torrent: %s", bangumi.official_title) return True else: logger.debug( "[Downloader] Torrent added before: %s", bangumi.official_title ) return False except Exception as e: logger.error( f"[Downloader] Failed to add torrent for {bangumi.official_title}: {e}" ) return False async def move_torrent(self, hashes, location): await self.client.move_torrent(hashes=hashes, new_location=location) # RSS Parts async def add_rss_feed(self, rss_link, item_path="Mikan_RSS"): await self.client.rss_add_feed(url=rss_link, item_path=item_path) async def remove_rss_feed(self, item_path): await self.client.rss_remove_item(item_path=item_path) async def get_rss_feed(self): return await self.client.rss_get_feeds() async def get_download_rules(self): return await self.client.get_download_rule() async def get_torrent_path(self, hashes): return await self.client.get_torrent_path(hashes) async def set_category(self, hashes, category): await self.client.set_category(hashes, category) async def remove_rule(self, rule_name): await self.client.remove_rule(rule_name) logger.info(f"[Downloader] Delete rule: {rule_name}") async def get_torrents_by_tag(self, tag: str) -> list[dict]: """Get all torrents with a specific tag.""" if hasattr(self.client, "get_torrents_by_tag"): return await self.client.get_torrents_by_tag(tag) return [] async def add_tag(self, torrent_hash: str, tag: str): """Add a tag to a torrent.""" await self.client.add_tag(torrent_hash, tag) logger.debug( "[Downloader] Added tag '%s' to torrent %s...", tag, torrent_hash[:8] )