From a030bca5d39ffd8d2e3cffcbfee9e3dc84fccee7 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 19 Jan 2026 20:23:45 +0000 Subject: [PATCH] Fix: Use info hash for clients (#495) Passes prowlarr's info hash to download client, existing behavior as fallback --- .../prowlarr/clients/__init__.py | 17 +++- .../prowlarr/clients/deluge.py | 11 ++- .../prowlarr/clients/nzbget.py | 10 ++- .../prowlarr/clients/qbittorrent.py | 12 ++- .../prowlarr/clients/rtorrent.py | 14 ++- .../prowlarr/clients/sabnzbd.py | 10 ++- .../prowlarr/clients/torrent_utils.py | 85 ++++++++++++++++++- .../prowlarr/clients/transmission.py | 12 ++- shelfmark/release_sources/prowlarr/handler.py | 2 + tests/prowlarr/test_clients.py | 14 +-- tests/prowlarr/test_failure_scenarios.py | 9 +- tests/prowlarr/test_integration_clients.py | 5 +- tests/prowlarr/test_qbittorrent_client.py | 51 +++++++++++ 13 files changed, 227 insertions(+), 25 deletions(-) diff --git a/shelfmark/release_sources/prowlarr/clients/__init__.py b/shelfmark/release_sources/prowlarr/clients/__init__.py index 510e84a..9e9b867 100644 --- a/shelfmark/release_sources/prowlarr/clients/__init__.py +++ b/shelfmark/release_sources/prowlarr/clients/__init__.py @@ -18,7 +18,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass from enum import Enum from functools import wraps -from typing import Callable, Dict, List, Optional, Tuple, Type, TypeVar, Union +from typing import Callable, Dict, List, Optional, Tuple, Type, TypeVar, Union, cast, Any import requests @@ -87,7 +87,9 @@ def with_retry( time.sleep(delay) # All retries exhausted - raise last_exception + if last_exception is None: + raise RuntimeError("Retry failed without exception") + raise cast(Exception, last_exception) return wrapper return decorator @@ -256,13 +258,22 @@ class DownloadClient(ABC): pass @abstractmethod - def add_download(self, url: str, name: str, category: Optional[str] = None) -> str: + def add_download( + self, + url: str, + name: str, + category: Optional[str] = None, + expected_hash: Optional[str] = None, + **kwargs: Any, + ) -> str: + """Add a download to the client. Args: url: Download URL (magnet link, .torrent URL, or NZB URL) name: Display name for the download category: Category/label for organization (None = client default) + expected_hash: Optional info_hash hint (torrents only) Returns: Client-specific download ID (hash for torrents, ID for NZBGet). diff --git a/shelfmark/release_sources/prowlarr/clients/deluge.py b/shelfmark/release_sources/prowlarr/clients/deluge.py index a678b50..a20d569 100644 --- a/shelfmark/release_sources/prowlarr/clients/deluge.py +++ b/shelfmark/release_sources/prowlarr/clients/deluge.py @@ -202,13 +202,20 @@ class DelugeClient(DownloadClient): self._connected = False return False, f"Connection failed: {str(e)}" - def add_download(self, url: str, name: str, category: Optional[str] = None) -> str: + def add_download( + self, + url: str, + name: str, + category: Optional[str] = None, + expected_hash: Optional[str] = None, + **kwargs, + ) -> str: try: self._ensure_connected() category_value = str(category or self._category) - torrent_info = extract_torrent_info(url) + torrent_info = extract_torrent_info(url, expected_hash=expected_hash) if not torrent_info.is_magnet and not torrent_info.torrent_data: raise Exception("Failed to fetch torrent file") diff --git a/shelfmark/release_sources/prowlarr/clients/nzbget.py b/shelfmark/release_sources/prowlarr/clients/nzbget.py index 4aec3fe..2b9a2bb 100644 --- a/shelfmark/release_sources/prowlarr/clients/nzbget.py +++ b/shelfmark/release_sources/prowlarr/clients/nzbget.py @@ -101,7 +101,14 @@ class NZBGetClient(DownloadClient): except Exception as e: return False, f"Connection failed: {str(e)}" - def add_download(self, url: str, name: str, category: Optional[str] = None) -> str: + def add_download( + self, + url: str, + name: str, + category: Optional[str] = None, + expected_hash: Optional[str] = None, + **kwargs, + ) -> str: """ Add NZB by URL. @@ -112,6 +119,7 @@ class NZBGetClient(DownloadClient): url: NZB URL (can be Prowlarr proxy URL) name: Display name for the download category: Category for organization (uses configured default if not specified) + expected_hash: Optional info_hash hint (unused) Returns: NZBGet download ID (NZBID). diff --git a/shelfmark/release_sources/prowlarr/clients/qbittorrent.py b/shelfmark/release_sources/prowlarr/clients/qbittorrent.py index aac6476..b7d4a6c 100644 --- a/shelfmark/release_sources/prowlarr/clients/qbittorrent.py +++ b/shelfmark/release_sources/prowlarr/clients/qbittorrent.py @@ -229,7 +229,14 @@ class QBittorrentClient(DownloadClient): except Exception as e: return False, f"Connection failed: {str(e)}" - def add_download(self, url: str, name: str, category: str | None = None) -> str: + def add_download( + self, + url: str, + name: str, + category: str | None = None, + expected_hash: str | None = None, + **kwargs, + ) -> str: """ Add torrent by URL (magnet or .torrent). @@ -237,6 +244,7 @@ class QBittorrentClient(DownloadClient): url: Magnet link or .torrent URL name: Display name for the torrent category: Category for organization (uses configured default if not specified) + expected_hash: Optional info_hash hint (from Prowlarr) Returns: Torrent hash (info_hash). @@ -257,7 +265,7 @@ class QBittorrentClient(DownloadClient): if "Conflict" not in type(e).__name__ and "409" not in str(e): logger.debug(f"Could not create category '{category}': {type(e).__name__}: {e}") - torrent_info = extract_torrent_info(url) + torrent_info = extract_torrent_info(url, expected_hash=expected_hash) expected_hash = torrent_info.info_hash torrent_data = torrent_info.torrent_data diff --git a/shelfmark/release_sources/prowlarr/clients/rtorrent.py b/shelfmark/release_sources/prowlarr/clients/rtorrent.py index f1231fc..02ff7f8 100644 --- a/shelfmark/release_sources/prowlarr/clients/rtorrent.py +++ b/shelfmark/release_sources/prowlarr/clients/rtorrent.py @@ -69,7 +69,14 @@ class RTorrentClient(DownloadClient): except Exception as e: return False, f"Connection failed: {str(e)}" - def add_download(self, url: str, name: str, category: Optional[str] = None) -> str: + def add_download( + self, + url: str, + name: str, + category: Optional[str] = None, + expected_hash: Optional[str] = None, + **kwargs, + ) -> str: """ Add torrent by URL (magnet or .torrent). @@ -77,6 +84,7 @@ class RTorrentClient(DownloadClient): url: Magnet link or .torrent URL name: Display name for the torrent category: Category for organization (uses configured label if not specified) + expected_hash: Optional info_hash hint (from Prowlarr) Returns: Torrent hash (info_hash). @@ -85,7 +93,7 @@ class RTorrentClient(DownloadClient): Exception: If adding fails. """ try: - torrent_info = extract_torrent_info(url) + torrent_info = extract_torrent_info(url, expected_hash=expected_hash) commands = [] @@ -109,7 +117,7 @@ class RTorrentClient(DownloadClient): add_url = torrent_info.magnet_url or url self._rpc.load.start("", add_url, ";".join(commands)) - torrent_hash = torrent_info.info_hash + torrent_hash = torrent_info.info_hash or expected_hash if not torrent_hash: raise Exception("Could not determine torrent hash from URL") diff --git a/shelfmark/release_sources/prowlarr/clients/sabnzbd.py b/shelfmark/release_sources/prowlarr/clients/sabnzbd.py index 03c90c5..969f44c 100644 --- a/shelfmark/release_sources/prowlarr/clients/sabnzbd.py +++ b/shelfmark/release_sources/prowlarr/clients/sabnzbd.py @@ -172,7 +172,14 @@ class SABnzbdClient(DownloadClient): except Exception as e: return False, f"Connection failed: {str(e)}" - def add_download(self, url: str, name: str, category: Optional[str] = None) -> str: + def add_download( + self, + url: str, + name: str, + category: Optional[str] = None, + expected_hash: Optional[str] = None, + **kwargs, + ) -> str: """ Add NZB by URL. @@ -180,6 +187,7 @@ class SABnzbdClient(DownloadClient): url: NZB URL (can be Prowlarr proxy URL) name: Display name for the download category: Category for organization (uses configured default if not specified) + expected_hash: Optional info_hash hint (unused) Returns: SABnzbd nzo_id. diff --git a/shelfmark/release_sources/prowlarr/clients/torrent_utils.py b/shelfmark/release_sources/prowlarr/clients/torrent_utils.py index 9160846..ce2714f 100644 --- a/shelfmark/release_sources/prowlarr/clients/torrent_utils.py +++ b/shelfmark/release_sources/prowlarr/clients/torrent_utils.py @@ -31,8 +31,23 @@ class TorrentInfo: magnet_url: Optional[str] = None """The actual magnet URL, if available.""" + def with_info_hash(self, info_hash: Optional[str]) -> "TorrentInfo": + """Return a copy with the info_hash replaced when provided.""" + if info_hash: + return TorrentInfo( + info_hash=info_hash, + torrent_data=self.torrent_data, + is_magnet=self.is_magnet, + magnet_url=self.magnet_url, + ) + return self -def extract_torrent_info(url: str, fetch_torrent: bool = True) -> TorrentInfo: + +def extract_torrent_info( + url: str, + fetch_torrent: bool = True, + expected_hash: Optional[str] = None, +) -> TorrentInfo: """Extract info_hash from magnet link or .torrent URL. Notes: @@ -48,12 +63,80 @@ def extract_torrent_info(url: str, fetch_torrent: bool = True) -> TorrentInfo: # Try to extract hash from magnet URL if is_magnet: info_hash = extract_hash_from_magnet(url) + if not info_hash and expected_hash: + info_hash = expected_hash return TorrentInfo(info_hash=info_hash, torrent_data=None, is_magnet=True, magnet_url=url) # Not a magnet - try to fetch and parse the .torrent file + if expected_hash: + return TorrentInfo(info_hash=expected_hash, torrent_data=None, is_magnet=False) + if not fetch_torrent: return TorrentInfo(info_hash=None, torrent_data=None, is_magnet=False) + headers: dict[str, str] = {} + api_key = str(config.get("PROWLARR_API_KEY", "") or "").strip() + if api_key: + headers["X-Api-Key"] = api_key + + def resolve_url(current: str, location: str) -> str: + if not location: + return current + # Support relative redirect locations + return urljoin(current, location) + + try: + logger.debug(f"Fetching torrent file from: {url[:80]}...") + + # Use allow_redirects=False to handle magnet link redirects manually + # Some indexers redirect download URLs to magnet links + resp = requests.get(url, timeout=30, allow_redirects=False, headers=headers) + + # Check if this is a redirect to a magnet link + if resp.status_code in (301, 302, 303, 307, 308): + redirect_url = resolve_url(url, resp.headers.get("Location", "")) + if redirect_url.startswith("magnet:"): + logger.debug("Download URL redirected to magnet link") + info_hash = extract_hash_from_magnet(redirect_url) + if not info_hash and expected_hash: + info_hash = expected_hash + return TorrentInfo( + info_hash=info_hash, torrent_data=None, is_magnet=True, magnet_url=redirect_url + ) + # Not a magnet redirect, follow it manually + logger.debug(f"Following redirect to: {redirect_url[:80]}...") + resp = requests.get(redirect_url, timeout=30, headers=headers) + + resp.raise_for_status() + torrent_data = resp.content + + # Check if response is actually a magnet link (text response) + # Some indexers return magnet links as plain text instead of redirecting + if len(torrent_data) < 2000: # Magnet links are typically short + try: + text_content = torrent_data.decode("utf-8", errors="ignore").strip() + if text_content.startswith("magnet:"): + logger.debug("Download URL returned magnet link as response body") + info_hash = extract_hash_from_magnet(text_content) + if not info_hash and expected_hash: + info_hash = expected_hash + return TorrentInfo( + info_hash=info_hash, torrent_data=None, is_magnet=True, magnet_url=text_content + ) + except Exception: + pass # Not text, continue with torrent parsing + + info_hash = extract_info_hash_from_torrent(torrent_data) + if info_hash: + logger.debug(f"Extracted hash from torrent file: {info_hash}") + else: + logger.warning("Could not extract hash from torrent file") + return TorrentInfo(info_hash=info_hash, torrent_data=torrent_data, is_magnet=False) + except Exception as e: + logger.debug(f"Could not fetch torrent file: {e}") + return TorrentInfo(info_hash=None, torrent_data=None, is_magnet=False) + + headers: dict[str, str] = {} api_key = str(config.get("PROWLARR_API_KEY", "") or "").strip() if api_key: diff --git a/shelfmark/release_sources/prowlarr/clients/transmission.py b/shelfmark/release_sources/prowlarr/clients/transmission.py index f115f14..1992740 100644 --- a/shelfmark/release_sources/prowlarr/clients/transmission.py +++ b/shelfmark/release_sources/prowlarr/clients/transmission.py @@ -73,7 +73,14 @@ class TransmissionClient(DownloadClient): except Exception as e: return False, f"Connection failed: {str(e)}" - def add_download(self, url: str, name: str, category: Optional[str] = None) -> str: + def add_download( + self, + url: str, + name: str, + category: Optional[str] = None, + expected_hash: Optional[str] = None, + **kwargs, + ) -> str: """ Add torrent by URL (magnet or .torrent). @@ -81,6 +88,7 @@ class TransmissionClient(DownloadClient): url: Magnet link or .torrent URL name: Display name for the torrent category: Category for organization (uses configured default if not specified) + expected_hash: Optional info_hash hint (from Prowlarr) Returns: Torrent hash (info_hash). @@ -91,7 +99,7 @@ class TransmissionClient(DownloadClient): try: resolved_category = category or self._category or "" - torrent_info = extract_torrent_info(url) + torrent_info = extract_torrent_info(url, expected_hash=expected_hash) if torrent_info.torrent_data: torrent = self._client.add_torrent( diff --git a/shelfmark/release_sources/prowlarr/handler.py b/shelfmark/release_sources/prowlarr/handler.py index 2679cd1..8ac47de 100644 --- a/shelfmark/release_sources/prowlarr/handler.py +++ b/shelfmark/release_sources/prowlarr/handler.py @@ -300,10 +300,12 @@ class ProwlarrHandler(DownloadHandler): try: release_name = prowlarr_result.get("title") or task.title or "Unknown" category = self._get_category_for_task(client, task) + expected_hash = str(prowlarr_result.get("infoHash") or "").strip() or None download_id = client.add_download( url=download_url, name=release_name, category=category, + expected_hash=expected_hash, ) except Exception as e: logger.error(f"Failed to add to {client.name}: {e}") diff --git a/tests/prowlarr/test_clients.py b/tests/prowlarr/test_clients.py index 8a33eaf..dc05dc0 100644 --- a/tests/prowlarr/test_clients.py +++ b/tests/prowlarr/test_clients.py @@ -151,7 +151,7 @@ class TestDownloadStatus: file_path=None, ) with pytest.raises(AttributeError): - status.progress = 75.0 + setattr(status, "progress", 75.0) def test_download_status_state_value_with_unknown_string(self): """Test state_value with an unknown state string.""" @@ -212,7 +212,7 @@ class TestClientRegistry: def test_connection(self): return True, "OK" - def add_download(self, url, name, category="test"): + def add_download(self, url, name, category=None, expected_hash=None, **kwargs): return "id" def get_status(self, download_id): @@ -245,7 +245,7 @@ class TestClientRegistry: def test_connection(self): return True, "OK" - def add_download(self, url, name, category="test"): + def add_download(self, url, name, category=None, expected_hash=None, **kwargs): return "id" def get_status(self, download_id): @@ -269,7 +269,7 @@ class TestClientRegistry: def test_connection(self): return True, "OK" - def add_download(self, url, name, category="test"): + def add_download(self, url, name, category=None, expected_hash=None, **kwargs): return "id" def get_status(self, download_id): @@ -301,7 +301,7 @@ class TestClientRegistry: def test_connection(self): return True, "OK" - def add_download(self, url, name, category="test"): + def add_download(self, url, name, category=None, expected_hash=None, **kwargs): return "id" def get_status(self, download_id): @@ -325,7 +325,7 @@ class TestClientRegistry: def test_connection(self): return True, "OK" - def add_download(self, url, name, category="test"): + def add_download(self, url, name, category=None, expected_hash=None, **kwargs): return "id" def get_status(self, download_id): @@ -377,7 +377,7 @@ class TestDownloadClientInterface: def test_connection(self): return True, "OK" - def add_download(self, url, name, category="test"): + def add_download(self, url, name, category=None, expected_hash=None, **kwargs): return "id" def get_status(self, download_id): diff --git a/tests/prowlarr/test_failure_scenarios.py b/tests/prowlarr/test_failure_scenarios.py index e04659b..b137c7c 100644 --- a/tests/prowlarr/test_failure_scenarios.py +++ b/tests/prowlarr/test_failure_scenarios.py @@ -82,7 +82,14 @@ class MockClient(DownloadClient): def test_connection(self) -> Tuple[bool, str]: return True, "Mock client connected" - def add_download(self, url: str, name: str, category: str = "cwabd") -> str: + def add_download( + self, + url: str, + name: str, + category: Optional[str] = None, + expected_hash: Optional[str] = None, + **kwargs, + ) -> str: if self.add_download_error: raise self.add_download_error download_id = f"mock-{len(self.downloads)}" diff --git a/tests/prowlarr/test_integration_clients.py b/tests/prowlarr/test_integration_clients.py index f800783..44d6aea 100644 --- a/tests/prowlarr/test_integration_clients.py +++ b/tests/prowlarr/test_integration_clients.py @@ -395,8 +395,9 @@ class TestQBittorrentIntegration: assert 0 <= status.progress <= 100 - valid_states = {"downloading", "complete", "error", "seeding", "paused", "queued", "fetching_metadata", "stalled"} - assert status.state.value in valid_states + valid_states = {"downloading", "complete", "error", "seeding", "paused", "queued", "fetching_metadata", "stalled", "checking"} + state_value = status.state.value if hasattr(status.state, "value") else status.state + assert state_value in valid_states assert isinstance(status.complete, bool) finally: diff --git a/tests/prowlarr/test_qbittorrent_client.py b/tests/prowlarr/test_qbittorrent_client.py index c82c266..bc15af7 100644 --- a/tests/prowlarr/test_qbittorrent_client.py +++ b/tests/prowlarr/test_qbittorrent_client.py @@ -10,6 +10,7 @@ from unittest.mock import MagicMock, patch import pytest from shelfmark.release_sources.prowlarr.clients import DownloadStatus +from shelfmark.release_sources.prowlarr.clients.torrent_utils import TorrentInfo class MockTorrent: @@ -465,6 +466,56 @@ class TestQBittorrentClientAddDownload: assert result == "3b245504cf5f11bbdbe1201cea6a6bf45aee1bc0" assert mock_client_instance._session.get.call_count >= 1 + def test_add_download_uses_expected_hash_without_fetch(self, monkeypatch): + """Skip proxy fetch when expected hash is provided for URL torrents.""" + config_values = { + "QBITTORRENT_URL": "http://localhost:8080", + "QBITTORRENT_USERNAME": "admin", + "QBITTORRENT_PASSWORD": "password", + "QBITTORRENT_CATEGORY": "test", + } + monkeypatch.setattr( + "shelfmark.release_sources.prowlarr.clients.qbittorrent.config.get", + lambda key, default="": config_values.get(key, default), + ) + + expected_hash = "3b245504cf5f11bbdbe1201cea6a6bf45aee1bc0" + mock_torrent = MockTorrent(hash_val=expected_hash) + mock_client_instance = MagicMock() + mock_client_instance.torrents_add.return_value = "Ok." + mock_client_instance.torrents_info.return_value = [mock_torrent] + mock_client_instance._session.get.return_value = create_mock_session_response({}, status_code=200) + mock_client_class = MagicMock(return_value=mock_client_instance) + + with patch.dict('sys.modules', {'qbittorrentapi': MagicMock(Client=mock_client_class)}): + import importlib + import shelfmark.release_sources.prowlarr.clients.qbittorrent as qb_module + importlib.reload(qb_module) + + with patch( + "shelfmark.release_sources.prowlarr.clients.qbittorrent.extract_torrent_info", + autospec=True, + ) as mock_extract: + mock_extract.return_value = TorrentInfo( + info_hash=expected_hash, + torrent_data=None, + is_magnet=False, + magnet_url=None, + ) + + client = qb_module.QBittorrentClient() + result = client.add_download( + "http://example.com/test.torrent", + "Test Download", + expected_hash=expected_hash, + ) + + assert result == expected_hash + mock_extract.assert_called_once_with( + "http://example.com/test.torrent", + expected_hash=expected_hash, + ) + def test_add_download_creates_category(self, monkeypatch): """Test that add_download creates category if needed.""" config_values = {