Improve qBittorrent support (#6)

* updated announce URL finder to also pull from tracker list

* Added full support for fastresume files

.gitignore

@@ -137,5 +137,7 @@ scratchpad.md
 tmp/
 
 *.torrent
+*.fastresume
 !tests/support/files/*.torrent
+!tests/support/files/*.fastresume
 tests/support/files/example.torrent

src/filesystem.py

@@ -19,3 +19,7 @@ def list_files_of_extension(input_directory: str, extension: str = ".torrent") -
   return [
     os.path.join(input_directory, filename) for filename in os.listdir(input_directory) if filename.endswith(extension)
   ]
+
+
+def replace_extension(filepath: str, new_extension: str) -> str:
+  return os.path.splitext(filepath)[0] + new_extension

src/injection.py

@@ -4,7 +4,7 @@
 from .errors import TorrentInjectionError
 from .clients.deluge import Deluge
 from .config import Config
-from .parser import calculate_infohash, get_name, get_torrent_data
+from .parser import calculate_infohash, get_name, get_bencoded_data
 
 
 class Injection:
@@ -18,7 +18,7 @@ def setup(self):
     return self
 
   def inject_torrent(self, source_torrent_filepath, new_torrent_filepath, new_tracker):
-    source_torrent_data = get_torrent_data(source_torrent_filepath)
+    source_torrent_data = get_bencoded_data(source_torrent_filepath)
     source_torrent_file_or_dir = self.__determine_torrent_data_location(source_torrent_data)
     output_location = self.__determine_output_location(source_torrent_file_or_dir, new_tracker)
     self.__link_files_to_output_location(source_torrent_file_or_dir, output_location)

src/parser.py

@@ -3,6 +3,7 @@
 import bencoder
 from hashlib import sha1
 
+from .utils import flatten
 from .trackers import RedTracker, OpsTracker
 
 
@@ -15,35 +16,40 @@ def is_valid_infohash(infohash: str) -> bool:
     return False
 
 
-def get_source(torrent_data: dict) -> bytes:
+def get_source(torrent_data: dict) -> bytes | None:
   try:
     return torrent_data[b"info"][b"source"]
   except KeyError:
     return None
 
 
-def get_name(torrent_data: dict) -> bytes:
+def get_name(torrent_data: dict) -> bytes | None:
   try:
     return torrent_data[b"info"][b"name"]
   except KeyError:
     return None
 
 
-def get_announce_url(torrent_data: dict) -> bytes:
-  try:
-    return torrent_data[b"announce"]
-  except KeyError:
-    return None
+def get_announce_url(torrent_data: dict) -> list[bytes] | None:
+  from_announce = torrent_data.get(b"announce")
+  if from_announce:
+    return from_announce if isinstance(from_announce, list) else [from_announce]
+
+  from_trackers = torrent_data.get(b"trackers")
+  if from_trackers:
+    return flatten(from_trackers)
+
+  return None
 
 
 def get_origin_tracker(torrent_data: dict) -> RedTracker | OpsTracker | None:
   source = get_source(torrent_data) or b""
-  announce_url = get_announce_url(torrent_data) or b""
+  announce_url = get_announce_url(torrent_data) or []
 
-  if source in RedTracker.source_flags_for_search() or RedTracker.announce_url() in announce_url:
+  if source in RedTracker.source_flags_for_search() or any(RedTracker.announce_url() in url for url in announce_url):
     return RedTracker
 
-  if source in OpsTracker.source_flags_for_search() or OpsTracker.announce_url() in announce_url:
+  if source in OpsTracker.source_flags_for_search() or any(OpsTracker.announce_url() in url for url in announce_url):
     return OpsTracker
 
   return None
@@ -60,7 +66,7 @@ def recalculate_hash_for_new_source(torrent_data: dict, new_source: (bytes | str
   return calculate_infohash(torrent_data)
 
 
-def get_torrent_data(filename: str) -> dict:
+def get_bencoded_data(filename: str) -> dict:
   try:
     with open(filename, "rb") as f:
       data = bencoder.decode(f.read())
@@ -69,7 +75,7 @@ def get_torrent_data(filename: str) -> dict:
     return None
 
 
-def save_torrent_data(filepath: str, torrent_data: dict) -> str:
+def save_bencoded_data(filepath: str, torrent_data: dict) -> str:
   parent_dir = os.path.dirname(filepath)
   os.makedirs(parent_dir, exist_ok=True)
 

src/scanner.py

@@ -4,7 +4,7 @@
 from .filesystem import mkdir_p, list_files_of_extension, assert_path_exists
 from .progress import Progress
 from .torrent import generate_new_torrent_from_file
-from .parser import get_torrent_data, calculate_infohash
+from .parser import get_bencoded_data, calculate_infohash
 from .errors import TorrentDecodingError, UnknownTrackerError, TorrentNotFoundError, TorrentAlreadyExistsError
 from .injection import Injection
 
@@ -134,7 +134,7 @@ def __collect_infohashes_from_files(files: list[str]) -> dict:
   infohash_dict = {}
 
   for filename in files:
-    torrent_data = get_torrent_data(filename)
+    torrent_data = get_bencoded_data(filename)
 
     if torrent_data:
       infohash = calculate_infohash(torrent_data)

src/torrent.py

@@ -5,11 +5,12 @@
 from .api import RedAPI, OpsAPI
 from .trackers import RedTracker, OpsTracker
 from .errors import TorrentDecodingError, UnknownTrackerError, TorrentNotFoundError, TorrentAlreadyExistsError
+from .filesystem import replace_extension
 from .parser import (
-  get_torrent_data,
+  get_bencoded_data,
   get_origin_tracker,
   recalculate_hash_for_new_source,
-  save_torrent_data,
+  save_bencoded_data,
 )
 
 
@@ -41,7 +42,7 @@ def generate_new_torrent_from_file(
     `Exception`: if an unknown error occurs.
   """
 
-  source_torrent_data, source_tracker = __get_torrent_data_and_tracker(source_torrent_path)
+  source_torrent_data, source_tracker = __get_bencoded_data_and_tracker(source_torrent_path)
   new_torrent_data = copy.deepcopy(source_torrent_data)
   new_tracker = source_tracker.reciprocal_tracker()
   new_tracker_api = __get_reciprocal_tracker_api(new_tracker, red_api, ops_api)
@@ -70,7 +71,7 @@ def generate_new_torrent_from_file(
         new_torrent_data[b"announce"] = new_tracker_api.announce_url.encode()
         new_torrent_data[b"comment"] = __generate_torrent_url(new_tracker_api.site_url, torrent_id).encode()
 
-        return (new_tracker, save_torrent_data(new_torrent_filepath, new_torrent_data))
+        return (new_tracker, save_bencoded_data(new_torrent_filepath, new_torrent_data))
     elif api_response["error"] in ("bad hash parameter", "bad parameters"):
       raise TorrentNotFoundError(f"Torrent could not be found on {new_tracker.site_shortname()}")
     else:
@@ -109,13 +110,23 @@ def __generate_torrent_url(site_url: str, torrent_id: str) -> str:
   return f"{site_url}/torrents.php?torrentid={torrent_id}"
 
 
-def __get_torrent_data_and_tracker(torrent_path):
-  source_torrent_data = get_torrent_data(torrent_path)
+def __get_bencoded_data_and_tracker(torrent_path):
+  # The fastresume stuff is to support qBittorrent since it doesn't store
+  # announce URLs in the torrent file IFF we're taking the file from `BT_backup`.
+  #
+  # qbit stores that information in a sidecar file that has the exact same name
+  # as the torrent file but with a `.fastresume` extension instead. It's also stored
+  # in a list of lists called `trackers` in this `.fastresume` file instead of `announce`.
+  fastresume_path = replace_extension(torrent_path, ".fastresume")
+  source_torrent_data = get_bencoded_data(torrent_path)
+  fastresume_data = get_bencoded_data(fastresume_path)
 
   if not source_torrent_data:
     raise TorrentDecodingError("Error decoding torrent file")
 
-  source_tracker = get_origin_tracker(source_torrent_data)
+  torrent_tracker = get_origin_tracker(source_torrent_data)
+  fastresume_tracker = get_origin_tracker(fastresume_data) if fastresume_data else None
+  source_tracker = torrent_tracker or fastresume_tracker
 
   if not source_tracker:
     raise UnknownTrackerError("Torrent not from OPS or RED based on source or announce URL")

src/utils.py

@@ -0,0 +1,4 @@
+def flatten(arg):
+  if not isinstance(arg, list):
+    return [arg]
+  return [x for sub in arg for x in flatten(sub)]

tests/test_filesystem.py

@@ -3,8 +3,7 @@
 
 from .helpers import SetupTeardown
 
-from src.filesystem import mkdir_p, assert_path_exists, list_files_of_extension
-# from src.errors import ConfigKeyError
+from src.filesystem import mkdir_p, assert_path_exists, list_files_of_extension, replace_extension
 
 
 class TestMkdirP(SetupTeardown):
@@ -63,3 +62,12 @@ def test_returns_empty_list_when_no_files_found(self):
     files = list_files_of_extension(input_directory, ".txt")
 
     assert len(files) == 0
+
+
+class TestReplaceExtension(SetupTeardown):
+  def test_replaces_extension(self):
+    filepath = "tests/support/files/test.torrent"
+
+    new_filepath = replace_extension(filepath, ".json")
+
+    assert new_filepath == "tests/support/files/test.json"

tests/test_parser.py

@@ -7,11 +7,11 @@
   is_valid_infohash,
   get_source,
   get_name,
-  get_torrent_data,
+  get_bencoded_data,
   get_announce_url,
   get_origin_tracker,
   recalculate_hash_for_new_source,
-  save_torrent_data,
+  save_bencoded_data,
   calculate_infohash,
 )
 
@@ -44,8 +44,14 @@ def test_returns_none_if_absent(self):
 
 
 class TestGetAnnounceUrl(SetupTeardown):
-  def test_returns_url_if_present(self):
-    assert get_announce_url({b"announce": b"https://foo.bar"}) == b"https://foo.bar"
+  def test_returns_url_if_present_in_announce(self):
+    assert get_announce_url({b"announce": b"https://foo.bar"}) == [b"https://foo.bar"]
+
+  def test_returns_url_if_present_in_trackers(self):
+    assert get_announce_url({b"trackers": [[b"https://foo.bar"], b"https://baz.qux"]}) == [
+      b"https://foo.bar",
+      b"https://baz.qux",
+    ]
 
   def test_returns_none_if_absent(self):
     assert get_announce_url({}) is None
@@ -65,6 +71,12 @@ def test_returns_red_based_on_announce(self):
   def test_returns_ops_based_on_announce(self):
     assert get_origin_tracker({b"announce": b"https://home.opsfet.ch/123abc"}) == OpsTracker
 
+  def test_returns_red_based_on_trackers(self):
+    assert get_origin_tracker({b"trackers": [[b"https://flacsfor.me/123abc"], b"https://baz.qux"]}) == RedTracker
+
+  def test_returns_ops_based_on_trackers(self):
+    assert get_origin_tracker({b"trackers": [[b"https://home.opsfet.ch/123abc"], b"https://baz.qux"]}) == OpsTracker
+
   def test_returns_none_if_no_match(self):
     assert get_origin_tracker({}) is None
     assert get_origin_tracker({b"info": {b"source": b"FOO"}}) is None
@@ -99,23 +111,23 @@ def test_doesnt_mutate_original_dict(self):
 
 class TestGetTorrentData(SetupTeardown):
   def test_returns_torrent_data(self):
-    result = get_torrent_data(get_torrent_path("no_source"))
+    result = get_bencoded_data(get_torrent_path("no_source"))
 
     assert isinstance(result, dict)
     assert b"info" in result
 
   def test_returns_none_on_error(self):
-    result = get_torrent_data(get_torrent_path("broken"))
+    result = get_bencoded_data(get_torrent_path("broken"))
 
     assert result is None
 
 
 class TestSaveTorrentData(SetupTeardown):
   def test_saves_torrent_data(self):
     torrent_data = {b"info": {b"source": b"RED"}}
-    filename = "/tmp/test_save_torrent_data.torrent"
+    filename = "/tmp/test_save_bencoded_data.torrent"
 
-    save_torrent_data(filename, torrent_data)
+    save_bencoded_data(filename, torrent_data)
 
     with open(filename, "rb") as f:
       result = f.read()
@@ -126,19 +138,19 @@ def test_saves_torrent_data(self):
 
   def test_returns_filename(self):
     torrent_data = {b"info": {b"source": b"RED"}}
-    filename = "/tmp/test_save_torrent_data.torrent"
+    filename = "/tmp/test_save_bencoded_data.torrent"
 
-    result = save_torrent_data(filename, torrent_data)
+    result = save_bencoded_data(filename, torrent_data)
 
     assert result == filename
 
     os.remove(filename)
 
   def test_creates_parent_directory(self):
     torrent_data = {b"info": {b"source": b"RED"}}
-    filename = "/tmp/output/foo/test_save_torrent_data.torrent"
+    filename = "/tmp/output/foo/test_save_bencoded_data.torrent"
 
-    save_torrent_data(filename, torrent_data)
+    save_bencoded_data(filename, torrent_data)
 
     assert os.path.exists("/tmp/output/foo")
 

tests/test_torrent.py

@@ -6,7 +6,7 @@
 from .helpers import get_torrent_path, SetupTeardown
 
 from src.trackers import RedTracker
-from src.parser import get_torrent_data
+from src.parser import get_bencoded_data
 from src.errors import TorrentAlreadyExistsError, TorrentDecodingError, UnknownTrackerError, TorrentNotFoundError
 from src.torrent import generate_new_torrent_from_file, generate_torrent_output_filepath
 
@@ -19,7 +19,7 @@ def test_saves_new_torrent_from_red_to_ops(self, red_api, ops_api):
 
       torrent_path = get_torrent_path("red_source")
       _, filepath = generate_new_torrent_from_file(torrent_path, "/tmp", red_api, ops_api)
-      parsed_torrent = get_torrent_data(filepath)
+      parsed_torrent = get_bencoded_data(filepath)
 
       assert os.path.isfile(filepath)
       assert parsed_torrent[b"announce"] == b"https://home.opsfet.ch/bar/announce"
@@ -35,7 +35,22 @@ def test_saves_new_torrent_from_ops_to_red(self, red_api, ops_api):
 
       torrent_path = get_torrent_path("ops_source")
       _, filepath = generate_new_torrent_from_file(torrent_path, "/tmp", red_api, ops_api)
-      parsed_torrent = get_torrent_data(filepath)
+      parsed_torrent = get_bencoded_data(filepath)
+
+      assert parsed_torrent[b"announce"] == b"https://flacsfor.me/bar/announce"
+      assert parsed_torrent[b"comment"] == b"https://redacted.ch/torrents.php?torrentid=123"
+      assert parsed_torrent[b"info"][b"source"] == b"RED"
+
+      os.remove(filepath)
+
+  def test_works_with_qbit_fastresume_files(self, red_api, ops_api):
+    with requests_mock.Mocker() as m:
+      m.get(re.compile("action=torrent"), json=self.TORRENT_SUCCESS_RESPONSE)
+      m.get(re.compile("action=index"), json=self.ANNOUNCE_SUCCESS_RESPONSE)
+
+      torrent_path = get_torrent_path("qbit_ops")
+      _, filepath = generate_new_torrent_from_file(torrent_path, "/tmp", red_api, ops_api)
+      parsed_torrent = get_bencoded_data(filepath)
 
       assert parsed_torrent[b"announce"] == b"https://flacsfor.me/bar/announce"
       assert parsed_torrent[b"comment"] == b"https://redacted.ch/torrents.php?torrentid=123"
@@ -50,7 +65,7 @@ def test_returns_new_tracker_instance_and_filepath(self, red_api, ops_api):
 
       torrent_path = get_torrent_path("ops_source")
       new_tracker, filepath = generate_new_torrent_from_file(torrent_path, "/tmp", red_api, ops_api)
-      get_torrent_data(filepath)
+      get_bencoded_data(filepath)
 
       assert os.path.isfile(filepath)
       assert new_tracker == RedTracker

tests/test_utils.py

@@ -0,0 +1,11 @@
+from .helpers import SetupTeardown
+
+from src.utils import flatten
+
+
+class TestFlatten(SetupTeardown):
+  def test_flattens_list(self):
+    assert flatten([1, [2, 3], 4]) == [1, 2, 3, 4]
+
+  def test_returns_already_flat_list(self):
+    assert flatten([1, 2, 3]) == [1, 2, 3]