add id search capabilities to jackett

media_manager/indexer/indexers/jackett.py

@@ -1,7 +1,9 @@
 import concurrent
 import concurrent.futures
 import logging
+import xml.etree.ElementTree as ET
 from concurrent.futures.thread import ThreadPoolExecutor
+from dataclasses import dataclass
 
 import requests
 
@@ -15,6 +17,21 @@ from media_manager.tv.schemas import Show
 log = logging.getLogger(__name__)
 
 
+@dataclass
+class IndexerInfo:
+    supports_tv_search: bool
+    supports_tv_search_tmdb: bool
+    supports_tv_search_imdb: bool
+    supports_tv_search_tvdb: bool
+    supports_tv_search_season: bool
+    supports_tv_search_episode: bool
+
+    supports_movie_search: bool
+    supports_movie_search_tmdb: bool
+    supports_movie_search_imdb: bool
+    supports_movie_search_tvdb: bool
+
+
 class Jackett(GenericIndexer, TorznabMixin):
     def __init__(self) -> None:
         """
@@ -31,11 +48,16 @@ class Jackett(GenericIndexer, TorznabMixin):
     def search(self, query: str, is_tv: bool) -> list[IndexerQueryResult]:
         log.debug("Searching for " + query)
 
+        params = {"q": query, "t": "tvsearch" if is_tv else "movie"}
+
+        return self.__search_jackett(params)
+
+    def __search_jackett(self, params: dict) -> list[IndexerQueryResult]:
         futures = []
         with ThreadPoolExecutor() as executor, requests.Session() as session:
             for indexer in self.indexers:
                 future = executor.submit(
-                    self.get_torrents_by_indexer, indexer, query, is_tv, session
+                    self.get_torrents_by_indexer, indexer, params, session
                 )
                 futures.append(future)
 
@@ -51,14 +73,103 @@ class Jackett(GenericIndexer, TorznabMixin):
 
         return responses
 
-    def get_torrents_by_indexer(
-        self, indexer: str, query: str, is_tv: bool, session: requests.Session
-    ) -> list[IndexerQueryResult]:
+    def __get_search_capabilities(
+        self, indexer: str, session: requests.Session
+    ) -> IndexerInfo:
         url = (
             self.url
-            + f"/api/v2.0/indexers/{indexer}/results/torznab/api?apikey={self.api_key}&t={'tvsearch' if is_tv else 'movie'}&q={query}"
+            + f"/api/v2.0/indexers/{indexer}/results/torznab/api?apikey={self.api_key}&t=caps"
         )
         response = session.get(url, timeout=self.timeout_seconds)
+        if response.status_code != 200:
+            msg = f"Cannot get search capabilities for Indexer {indexer}"
+            log.error(msg)
+            raise RuntimeError(msg)
+
+        xml = response.text
+        xml_tree = ET.fromstring(xml)  # noqa: S314  # trusted source, since it is user controlled
+        tv_search = xml_tree.find("./*/tv-search")
+        movie_search = xml_tree.find("./*/movie-search")
+        log.debug(tv_search.attrib)
+        log.debug(movie_search.attrib)
+
+        tv_search_capabilities = []
+        movie_search_capabilities = []
+        tv_search_available = (tv_search is not None) and (
+            tv_search.attrib["available"] == "yes"
+        )
+        movie_search_available = (movie_search is not None) and (
+            movie_search.attrib["available"] == "yes"
+        )
+
+        if tv_search_available:
+            tv_search_capabilities = tv_search.attrib["supportedParams"].split(",")
+
+        if movie_search_available:
+            movie_search_capabilities = movie_search.attrib["supportedParams"].split(
+                ","
+            )
+
+        return IndexerInfo(
+            supports_tv_search=tv_search_available,
+            supports_tv_search_imdb="tmdbid" in tv_search_capabilities,
+            supports_tv_search_tmdb="tmdbid" in tv_search_capabilities,
+            supports_tv_search_tvdb="tvdbid" in tv_search_capabilities,
+            supports_tv_search_season="season" in tv_search_capabilities,
+            supports_tv_search_episode="e" in tv_search_capabilities,
+            supports_movie_search=movie_search_available,
+            supports_movie_search_imdb="imdbid" in movie_search_capabilities,
+            supports_movie_search_tmdb="tmdbid" in movie_search_capabilities,
+            supports_movie_search_tvdb="tvdbid" in movie_search_capabilities,
+        )
+
+    def __get_optimal_query_parameters(
+        self, indexer: str, session: requests.Session, params: dict
+    ) -> dict[str, str]:
+        query_params = {"apikey": self.api_key, "t": params["t"]}
+
+        search_capabilities = self.__get_search_capabilities(
+            indexer=indexer, session=session
+        )
+        if params["t"] == "tvsearch":
+            if not search_capabilities.supports_tv_search:
+                msg = f"Indexer {indexer} does not support TV search"
+                raise RuntimeError(msg)
+            if search_capabilities.supports_tv_search_season and "season" in params:
+                query_params["season"] = params["season"]
+            if search_capabilities.supports_tv_search_episode and "e" in params:
+                query_params["e"] = params["e"]
+            if search_capabilities.supports_tv_search_imdb and "imdbid" in params:
+                query_params["imdbid"] = params["imdbid"]
+            elif search_capabilities.supports_tv_search_tvdb and "tvdbid" in params:
+                query_params["tvdbid"] = params["tvdbid"]
+            elif search_capabilities.supports_tv_search_tmdb and "tmdbid" in params:
+                query_params["tmdbid"] = params["tmdbid"]
+            else:
+                query_params["q"] = params["q"]
+        if params["t"] == "movie":
+            if not search_capabilities.supports_movie_search:
+                msg = f"Indexer {indexer} does not support Movie search"
+                raise RuntimeError(msg)
+            if search_capabilities.supports_movie_search_imdb and "imdbid" in params:
+                query_params["imdbid"] = params["imdbid"]
+            elif search_capabilities.supports_tv_search_tvdb and "tvdbid" in params:
+                query_params["tvdbid"] = params["tvdbid"]
+            elif search_capabilities.supports_tv_search_tmdb and "tmdbid" in params:
+                query_params["tmdbid"] = params["tmdbid"]
+            else:
+                query_params["q"] = params["q"]
+        return query_params
+
+    def get_torrents_by_indexer(
+        self, indexer: str, params: dict, session: requests.Session
+    ) -> list[IndexerQueryResult]:
+        url = f"{self.url}/api/v2.0/indexers/{indexer}/results/torznab/api"
+        query_params = self.__get_optimal_query_parameters(
+            indexer=indexer, session=session, params=params
+        )
+        response = session.get(url, timeout=self.timeout_seconds, params=query_params)
+        log.debug(f"Indexer {indexer} url: {response.url}")
 
         if response.status_code != 200:
             log.error(
@@ -75,8 +186,23 @@ class Jackett(GenericIndexer, TorznabMixin):
         self, query: str, show: Show, season_number: int
     ) -> list[IndexerQueryResult]:
         log.debug(f"Searching for season {season_number} of show {show.name}")
-        return self.search(query=query, is_tv=True)
+        params = {
+            "t": "tvsearch",
+            "season": season_number,
+            "q": query,
+        }
+        if show.imdb_id:
+            params["imdbid"] = show.imdb_id
+        params[show.metadata_provider + "id"] = show.external_id
+        return self.__search_jackett(params=params)
 
     def search_movie(self, query: str, movie: Movie) -> list[IndexerQueryResult]:
         log.debug(f"Searching for movie {movie.name}")
-        return self.search(query=query, is_tv=False)
+        params = {
+            "t": "movie",
+            "q": query,
+        }
+        if movie.imdb_id:
+            params["imdbid"] = movie.imdb_id
+        params[movie.metadata_provider + "id"] = movie.external_id
+        return self.__search_jackett(params=params)