diff --git a/hate_crack/api.py b/hate_crack/api.py index f40121b..34c94bf 100644 --- a/hate_crack/api.py +++ b/hate_crack/api.py @@ -2,10 +2,11 @@ import concurrent.futures import json import sys import os +import shutil +import tempfile import threading import time from queue import Queue -import shutil from typing import Callable, Optional, Tuple import requests # type: ignore[import-untyped] @@ -258,6 +259,7 @@ class TransmissionSession: self.startup_timeout = startup_timeout self.shutdown_timeout = shutdown_timeout self._cfg_dir = "" + self._watch_dir = "" self._port = 0 self._rpc = "" self._proc = None @@ -266,9 +268,10 @@ class TransmissionSession: def __enter__(self): import atexit import subprocess - import tempfile self._cfg_dir = tempfile.mkdtemp(prefix="hate_crack_transmission_") + self._watch_dir = os.path.join(self._cfg_dir, "watch") + os.makedirs(self._watch_dir, exist_ok=True) self._port = _pick_free_port() self._rpc = f"127.0.0.1:{self._port}" self._proc = subprocess.Popen( @@ -285,8 +288,11 @@ class TransmissionSession: "--download-dir", self.save_dir, "--no-portmap", - "--no-watch-dir", - ] + "--watch-dir", + self._watch_dir, + ], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, ) deadline = time.monotonic() + self.startup_timeout while time.monotonic() < deadline: @@ -339,32 +345,15 @@ class TransmissionSession: return None def add(self, torrent_path: str) -> int: - import re - import subprocess - before_ids = {e["id"] for e in self.list()} - result = subprocess.run( - [ - "transmission-remote", - self._rpc, - "--no-auth", - "-a", - torrent_path, - ], - capture_output=True, - text=True, - ) - out = result.stdout or "" - m = re.search(r"Added torrent.*\n.*ID:\s*(\d+)", out) - if m: - return int(m.group(1)) - m = re.search(r"torrent added\s*\(id\s+(\d+)\)", out, re.IGNORECASE) - if m: - return int(m.group(1)) - after_entries = self.list() - new_ids = [e["id"] for e in after_entries if e["id"] not in before_ids] - if new_ids: - return new_ids[0] + shutil.copy2(torrent_path, self._watch_dir) + deadline = time.monotonic() + 10.0 + while time.monotonic() < deadline: + after_entries = self.list() + new_ids = [e["id"] for e in after_entries if e["id"] not in before_ids] + if new_ids: + return new_ids[0] + time.sleep(0.5) raise RuntimeError(f"Failed to add torrent: {torrent_path}") def list(self) -> list: @@ -582,9 +571,9 @@ def get_hcat_potfile_args(): def cleanup_torrent_files(directory=None): - """Remove stray .torrent files from the wordlists directory on graceful exit.""" + """Remove stray .torrent files left in the system temp directory on graceful exit.""" if directory is None: - directory = get_hcat_wordlists_dir() + directory = tempfile.gettempdir() try: for name in os.listdir(directory): if name.endswith(".torrent"): @@ -791,18 +780,12 @@ def fetch_torrent_metadata(torrent_url, save_dir=None, wordlist_id=None): """Download the .torrent metadata file from Weakpass and return its local path. Returns the path to the saved .torrent file, or None on failure. + The .torrent file is stored in the system temp directory, not the wordlist dir. """ register_torrent_cleanup() - if not save_dir: - save_dir = get_hcat_wordlists_dir() - else: - save_dir = os.path.expanduser(save_dir) - if not os.path.isabs(save_dir): - save_dir = os.path.join( - os.path.dirname(os.path.abspath(__file__)), save_dir - ) - os.makedirs(save_dir, exist_ok=True) + torrent_dir = tempfile.gettempdir() + os.makedirs(torrent_dir, exist_ok=True) # Optionally include hashmob_api_key in headers if present headers = { "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36" @@ -899,7 +882,7 @@ def fetch_torrent_metadata(torrent_url, save_dir=None, wordlist_id=None): r2 = requests.get(torrent_link, headers=headers, stream=True) content_type = r2.headers.get("Content-Type", "") local_filename = os.path.join( - save_dir, filename if filename.endswith(".torrent") else filename + ".torrent" + torrent_dir, filename if filename.endswith(".torrent") else filename + ".torrent" ) if r2.status_code == 200 and not content_type.startswith("text/html"): with open(local_filename, "wb") as f: diff --git a/tests/test_api_downloads.py b/tests/test_api_downloads.py index ef27a5a..cb21d7e 100644 --- a/tests/test_api_downloads.py +++ b/tests/test_api_downloads.py @@ -140,48 +140,44 @@ class TestTransmissionSession: with pytest.raises(RuntimeError, match="Transmission daemon failed"): ts.__enter__() - def test_add_parses_id_from_stdout(self, tmp_path): + def test_add_copies_to_watch_dir_and_returns_new_id(self, tmp_path): ts = TransmissionSession(str(tmp_path)) ts._rpc = "127.0.0.1:9999" - result = MagicMock( - returncode=0, stdout="Added torrent foo.torrent\n ID: 7\n", stderr="" - ) - with patch("subprocess.run", return_value=result): - tid = ts.add("/tmp/foo.torrent") - assert tid == 7 - - def test_add_parses_lowercase_alt_format(self, tmp_path): - ts = TransmissionSession(str(tmp_path)) - ts._rpc = "127.0.0.1:9999" - result = MagicMock( - returncode=0, stdout="torrent added (id 42)\n", stderr="" - ) - with patch("subprocess.run", return_value=result): - tid = ts.add("/tmp/foo.torrent") - assert tid == 42 - - def test_add_falls_back_to_list(self, tmp_path): - ts = TransmissionSession(str(tmp_path)) - ts._rpc = "127.0.0.1:9999" - result = MagicMock(returncode=0, stdout="garbage output\n", stderr="") - # Before: IDs 3 and 5 exist. After: ID 7 appears as the newly added torrent. + ts._watch_dir = str(tmp_path / "watch") + # Before: IDs 3 and 5. After first poll: ID 7 appears. list_calls = iter([ - [{"id": 3}, {"id": 5}], # before snapshot - [{"id": 3}, {"id": 5}, {"id": 7}], # after snapshot + [{"id": 3}, {"id": 5}], + [{"id": 3}, {"id": 5}, {"id": 7}], ]) - with patch("subprocess.run", return_value=result), patch.object( - ts, "list", side_effect=list_calls - ): + with patch("shutil.copy2"), patch.object(ts, "list", side_effect=list_calls), \ + patch("time.sleep"), patch("time.monotonic", side_effect=[0.0, 1.0]): tid = ts.add("/tmp/foo.torrent") assert tid == 7 - def test_add_raises_when_list_empty(self, tmp_path): + def test_add_polls_until_daemon_picks_up_torrent(self, tmp_path): ts = TransmissionSession(str(tmp_path)) ts._rpc = "127.0.0.1:9999" - result = MagicMock(returncode=0, stdout="garbage\n", stderr="") - with patch("subprocess.run", return_value=result), patch.object( - ts, "list", return_value=[] - ): + ts._watch_dir = str(tmp_path / "watch") + # First two polls: no new ID. Third poll: ID 9 appears. + list_calls = iter([ + [{"id": 1}], + [{"id": 1}], + [{"id": 1}], + [{"id": 1}, {"id": 9}], + ]) + monotonic_vals = iter([0.0, 1.0, 2.0, 3.0, 4.0]) + with patch("shutil.copy2"), patch.object(ts, "list", side_effect=list_calls), \ + patch("time.sleep"), patch("time.monotonic", side_effect=monotonic_vals): + tid = ts.add("/tmp/foo.torrent") + assert tid == 9 + + def test_add_raises_on_timeout(self, tmp_path): + ts = TransmissionSession(str(tmp_path)) + ts._rpc = "127.0.0.1:9999" + ts._watch_dir = str(tmp_path / "watch") + # list never returns a new ID; monotonic jumps past the 10s deadline. + with patch("shutil.copy2"), patch.object(ts, "list", return_value=[{"id": 1}]), \ + patch("time.sleep"), patch("time.monotonic", side_effect=[0.0, 100.0]): with pytest.raises(RuntimeError): ts.add("/tmp/foo.torrent")