From e5e1107359fb8916d8ba6c4c05e340ed0f6b1b83 Mon Sep 17 00:00:00 2001 From: Justin Bollinger Date: Fri, 24 Jul 2026 17:49:24 -0400 Subject: [PATCH] fix(llm): add spinner during Ollama generation and cap wordlist sample MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Defect 1 — hcatOllama showed a plain print then blocked silently for up to 300 s while waiting for the LLM. A new generic context manager `hate_crack/progress.py::spinner()` now runs a daemon thread that repaints a single line with a frame + elapsed-seconds counter every ~120 ms. TTY guard ensures non-TTY/test environments get a single plain print instead of ANSI control characters. The line is erased with \033[2K\r on exit (normal and exceptional). Defect 2 — the wordlist path materialised every line in memory and built a prompt that could massively overflow ollamaNumCtx on large wordlists. The path now does a two-pass evenly-spaced sample: pass 1 counts usable lines, pass 2 stride-selects up to `ollamaMaxSampleLines` (default 500, new config key) spread across the full file. When no capping occurs the message reads "Loaded N passwords"; when capping occurs it reads "Sampled N of M passwords". Co-Authored-By: Claude --- config.json.example | 1 + hate_crack/main.py | 88 +++++++--- hate_crack/progress.py | 65 +++++++ tests/test_ollama_wordlist_sampling.py | 231 +++++++++++++++++++++++++ tests/test_progress_spinner.py | 108 ++++++++++++ 5 files changed, 472 insertions(+), 21 deletions(-) create mode 100644 hate_crack/progress.py create mode 100644 tests/test_ollama_wordlist_sampling.py create mode 100644 tests/test_progress_spinner.py diff --git a/config.json.example b/config.json.example index 9efabe3..86bd196 100644 --- a/config.json.example +++ b/config.json.example @@ -28,6 +28,7 @@ "ollamaModel": "qwen2.5:32b", "ollamaNumCtx": 2048, "ollamaTimeout": 300, + "ollamaMaxSampleLines": 500, "omenTrainingList": "rockyou.txt", "omenMaxCandidates": 50000000, "pcfgRuleset": "DEFAULT", diff --git a/hate_crack/main.py b/hate_crack/main.py index b801441..32e50e1 100755 --- a/hate_crack/main.py +++ b/hate_crack/main.py @@ -74,6 +74,7 @@ from hate_crack.cli import ( # noqa: E402 ) from hate_crack import attacks as _attacks # noqa: E402 from hate_crack import llm # noqa: E402 +from hate_crack.progress import spinner # noqa: E402 from hate_crack.menu import interactive_menu # noqa: E402 from hate_crack.username_detect import detect_username_hash_format # noqa: E402 @@ -453,6 +454,7 @@ ollamaUrl = "http://" + os.environ.get("OLLAMA_HOST", "localhost:11434") ollamaModel = config_parser.get("ollamaModel", "qwen2.5:32b") ollamaNumCtx = int(config_parser.get("ollamaNumCtx", 2048)) ollamaTimeout = float(config_parser.get("ollamaTimeout", 300)) +ollamaMaxSampleLines = int(config_parser.get("ollamaMaxSampleLines", 500)) omenTrainingList = config_parser.get("omenTrainingList", "rockyou.txt") omenMaxCandidates = int(config_parser.get("omenMaxCandidates", 1000000)) @@ -2046,23 +2048,67 @@ def hcatOllama(hcatHashType, hcatHashFile, mode, context_data): if not os.path.isfile(wordlist_path): print(f"Error: Wordlist not found: {wordlist_path}") return - lines = [] + + # Two-pass evenly-spaced sample: first count usable lines so we can + # stride-select across the whole file rather than taking a head slice. + # A head-only sample misses the pattern variation across large wordlists + # (e.g. rockyou.txt becomes more random further in). + def _usable(raw: str) -> str: + """Return the usable plaintext from a raw line, or empty string.""" + stripped = raw.strip() + if not stripped: + return "" + if ":" in stripped: + stripped = stripped.split(":", 1)[1] + return stripped + try: + total_usable = 0 with open(wordlist_path, "r", errors="ignore") as f: - for line in f: - stripped = line.strip() - if not stripped: - continue - # hash:password -> password - if ":" in stripped: - stripped = stripped.split(":", 1)[1] - if stripped: - lines.append(stripped) + for raw in f: + if _usable(raw): + total_usable += 1 except Exception as e: print(f"Error reading wordlist: {e}") return - print(f"Loaded {len(lines)} passwords from wordlist.") - gen_context = {"sample": "\n".join(lines)} + + cap = ollamaMaxSampleLines + if total_usable <= cap: + # No capping needed — collect all usable lines. + try: + sampled: list[str] = [] + with open(wordlist_path, "r", errors="ignore") as f: + for raw in f: + w = _usable(raw) + if w: + sampled.append(w) + except Exception as e: + print(f"Error reading wordlist: {e}") + return + print(f"Loaded {len(sampled):,} passwords from wordlist.") + else: + # Evenly-spaced stride across the file to cover its full pattern + # range without materialising all lines in memory. + stride = total_usable / cap + try: + sampled = [] + usable_idx = 0.0 + next_pick = stride / 2 # start near centre of first bucket + with open(wordlist_path, "r", errors="ignore") as f: + for raw in f: + w = _usable(raw) + if not w: + continue + if usable_idx >= next_pick and len(sampled) < cap: + sampled.append(w) + next_pick += stride + usable_idx += 1 + except Exception as e: + print(f"Error reading wordlist: {e}") + return + print(f"Sampled {len(sampled):,} of {total_usable:,} passwords from wordlist.") + + gen_context = {"sample": "\n".join(sampled)} elif mode == "target": gen_context = context_data else: @@ -2070,16 +2116,16 @@ def hcatOllama(hcatHashType, hcatHashFile, mode, context_data): return # Step B: generate candidates via the Atomic Agents module. - print(f"Generating password candidates via Ollama ({ollamaModel})...") try: - candidates = llm.generate_candidates( - ollamaUrl, - ollamaModel, - ollamaNumCtx, - mode, - gen_context, - timeout=ollamaTimeout, - ) + with spinner(f"Generating password candidates via Ollama ({ollamaModel})..."): + candidates = llm.generate_candidates( + ollamaUrl, + ollamaModel, + ollamaNumCtx, + mode, + gen_context, + timeout=ollamaTimeout, + ) except llm.LLMTimeoutError: print(f"Error: the Ollama request timed out after {ollamaTimeout:g} seconds.") print( diff --git a/hate_crack/progress.py b/hate_crack/progress.py new file mode 100644 index 0000000..3827533 --- /dev/null +++ b/hate_crack/progress.py @@ -0,0 +1,65 @@ +"""Generic terminal progress utilities for hate_crack. + +Provides a context manager that shows a live spinner with elapsed-seconds +counter while a blocking operation runs. Safe to use in non-TTY environments: +if stdout is not a TTY the message is printed once and no background thread is +started. +""" + +from __future__ import annotations + +import sys +import threading +import time +from collections.abc import Generator +from contextlib import contextmanager + +_SPINNER_FRAMES = ["|", "/", "-", "\\"] +_TICK_INTERVAL = 0.12 # seconds between repaints (~120 ms) + + +@contextmanager +def spinner(message: str) -> "Generator[None, None, None]": + """Context manager that shows *message* plus a live elapsed-seconds counter. + + While the body executes a daemon thread repaints a single terminal line + roughly every 120 ms showing: + + | Generating password candidates via Ollama (model)... 3s + + On exit (normal or exceptional) the spinner line is erased so subsequent + output starts on a clean line. + + TTY guard: if ``sys.stdout.isatty()`` is False the message is printed once + via ``print()`` and no thread is started, keeping piped output and the test + suite clean. + """ + if not sys.stdout.isatty(): + print(message) + yield + return + + stop_event = threading.Event() + start_time = time.monotonic() + + def _run() -> None: + frame_idx = 0 + while not stop_event.is_set(): + elapsed = int(time.monotonic() - start_time) + frame = _SPINNER_FRAMES[frame_idx % len(_SPINNER_FRAMES)] + line = f"\r{frame} {message} {elapsed}s" + sys.stdout.write(line) + sys.stdout.flush() + frame_idx += 1 + stop_event.wait(_TICK_INTERVAL) + + thread = threading.Thread(target=_run, daemon=True) + thread.start() + try: + yield + finally: + stop_event.set() + thread.join() + # Erase the spinner line so the next print starts clean. + sys.stdout.write("\033[2K\r") + sys.stdout.flush() diff --git a/tests/test_ollama_wordlist_sampling.py b/tests/test_ollama_wordlist_sampling.py new file mode 100644 index 0000000..305ab4c --- /dev/null +++ b/tests/test_ollama_wordlist_sampling.py @@ -0,0 +1,231 @@ +"""Tests for the capped / evenly-sampled wordlist path in hcatOllama. + +Covers: +- ollamaMaxSampleLines config default (500) +- Small wordlist (< cap): all lines included, "Loaded N" message +- Large wordlist (> cap): exactly `cap` lines sampled, "Sampled N of M" message +- Evenly-spaced sample covers the whole file (first and last entries present) +- hash:password splitting and blank-line skipping still work +- spinner is called with the right message regardless of TTY +""" + +from __future__ import annotations + +import os +from contextlib import contextmanager +from types import SimpleNamespace +from unittest import mock + +import pytest + +os.environ["HATE_CRACK_SKIP_INIT"] = "1" +from hate_crack import main as hc_main # noqa: E402 + +OLLAMA_URL = "http://localhost:11434" +MODEL = "test-model" + + +@contextmanager +def _ollama_globals(tmp_path, *, max_sample: int = 500): + rules_dir = str(tmp_path / "rules") + os.makedirs(rules_dir, exist_ok=True) + with ( + mock.patch.object(hc_main, "ollamaUrl", OLLAMA_URL), + mock.patch.object(hc_main, "ollamaModel", MODEL), + mock.patch.object(hc_main, "ollamaNumCtx", 2048), + mock.patch.object(hc_main, "ollamaMaxSampleLines", max_sample), + mock.patch.object(hc_main, "hcatBin", "/usr/bin/hashcat"), + mock.patch.object(hc_main, "hcatTuning", ""), + mock.patch.object(hc_main, "hcatPotfilePath", ""), + mock.patch.object(hc_main, "rulesDirectory", rules_dir), + mock.patch("hate_crack.main.generate_session_id", return_value="s"), + ): + yield + + +def _make_proc(rc: int = 0): + proc = mock.MagicMock() + proc.wait.return_value = rc + proc.communicate.return_value = (b"", b"") + proc.returncode = rc + return proc + + +@pytest.fixture +def env(tmp_path): + hash_file = tmp_path / "hashes.txt" + hash_file.touch() + return SimpleNamespace(tmp_path=tmp_path, hash_file=str(hash_file)) + + +# --------------------------------------------------------------------------- +# Config default +# --------------------------------------------------------------------------- + + +def test_ollama_max_sample_lines_default(): + """ollamaMaxSampleLines must default to 500.""" + assert hc_main.ollamaMaxSampleLines == 500 + + +# --------------------------------------------------------------------------- +# Small wordlist — no capping +# --------------------------------------------------------------------------- + + +def test_small_wordlist_loads_all(env, capsys): + """When total lines < cap, all are included and message says 'Loaded'.""" + wl = env.tmp_path / "small.txt" + wl.write_text("alpha\nbeta\ngamma\n") + + with _ollama_globals(env.tmp_path), mock.patch.object( + hc_main.llm, "generate_candidates", return_value=["x"] + ) as gen, mock.patch("subprocess.Popen", return_value=_make_proc()): + hc_main.hcatOllama("0", env.hash_file, "wordlist", str(wl)) + + captured = capsys.readouterr() + assert "Loaded 3" in captured.out + # All three words must appear in the sample passed to the LLM + sample = gen.call_args[0][4]["sample"] + assert "alpha" in sample + assert "beta" in sample + assert "gamma" in sample + + +def test_small_wordlist_no_sampled_message(env, capsys): + """'Sampled N of M' must NOT appear when no capping occurred.""" + wl = env.tmp_path / "small.txt" + wl.write_text("a\nb\nc\n") + + with _ollama_globals(env.tmp_path, max_sample=500), mock.patch.object( + hc_main.llm, "generate_candidates", return_value=["x"] + ), mock.patch("subprocess.Popen", return_value=_make_proc()): + hc_main.hcatOllama("0", env.hash_file, "wordlist", str(wl)) + + captured = capsys.readouterr() + assert "Sampled" not in captured.out + + +# --------------------------------------------------------------------------- +# Large wordlist — capping +# --------------------------------------------------------------------------- + + +def _make_large_wordlist(path, n: int) -> None: + """Write n unique numbered lines to path.""" + lines = "\n".join(f"word{i:06d}" for i in range(n)) + path.write_text(lines + "\n") + + +def test_large_wordlist_caps_to_max(env, capsys): + """When total > cap, exactly cap lines are sampled.""" + wl = env.tmp_path / "big.txt" + _make_large_wordlist(wl, 1000) + + captured_ctx: list[dict] = [] + + def _capture_gen(*args, **kwargs): + captured_ctx.append(args[4]) + return ["x"] + + with _ollama_globals(env.tmp_path, max_sample=50), mock.patch.object( + hc_main.llm, "generate_candidates", side_effect=_capture_gen + ), mock.patch("subprocess.Popen", return_value=_make_proc()): + hc_main.hcatOllama("0", env.hash_file, "wordlist", str(wl)) + + sample_lines = captured_ctx[0]["sample"].splitlines() + assert len(sample_lines) == 50 + + captured = capsys.readouterr() + assert "Sampled 50 of 1,000" in captured.out + + +def test_large_wordlist_covers_full_range(env): + """Evenly-spaced sample must contain entries from both the start and end of the file.""" + wl = env.tmp_path / "range.txt" + _make_large_wordlist(wl, 1000) + + captured_ctx: list[dict] = [] + + def _capture_gen(*args, **kwargs): + captured_ctx.append(args[4]) + return ["x"] + + # Use a small cap (10) over 1000 lines so the stride is clearly 100. + with _ollama_globals(env.tmp_path, max_sample=10), mock.patch.object( + hc_main.llm, "generate_candidates", side_effect=_capture_gen + ), mock.patch("subprocess.Popen", return_value=_make_proc()): + hc_main.hcatOllama("0", env.hash_file, "wordlist", str(wl)) + + sample_lines = captured_ctx[0]["sample"].splitlines() + # The sampled words should come from across the file, not just the head. + # word000000..word000099 are the first 100; word000900..word000999 are the last 100. + indices = [int(w.replace("word", "")) for w in sample_lines] + assert min(indices) < 100, "sample should include entries from the start of the file" + assert max(indices) >= 900, "sample should include entries from the end of the file" + + +# --------------------------------------------------------------------------- +# hash:password splitting and blank-line skipping +# --------------------------------------------------------------------------- + + +def test_wordlist_sampling_strips_hash_prefix(env): + """hash:password lines must contribute only the plaintext in capped mode.""" + wl = env.tmp_path / "dump.txt" + # 20 lines: half with colon prefix, half plain; more than max_sample=5 + lines = [f"hash{i}:plain{i:02d}" if i % 2 == 0 else f"plain{i:02d}" for i in range(20)] + wl.write_text("\n".join(lines) + "\n") + + captured_ctx: list[dict] = [] + + def _capture_gen(*args, **kwargs): + captured_ctx.append(args[4]) + return ["x"] + + with _ollama_globals(env.tmp_path, max_sample=5), mock.patch.object( + hc_main.llm, "generate_candidates", side_effect=_capture_gen + ), mock.patch("subprocess.Popen", return_value=_make_proc()): + hc_main.hcatOllama("0", env.hash_file, "wordlist", str(wl)) + + sample = captured_ctx[0]["sample"] + # No hash prefix should appear in the sample + assert "hash" not in sample + + +def test_wordlist_sampling_skips_blank_lines(env, capsys): + """Blank lines must not count toward total or be included in the sample.""" + wl = env.tmp_path / "blanks.txt" + # 3 real words surrounded by blank lines + wl.write_text("\nalpha\n\nbeta\n\ngamma\n\n") + + with _ollama_globals(env.tmp_path, max_sample=500), mock.patch.object( + hc_main.llm, "generate_candidates", return_value=["x"] + ) as gen, mock.patch("subprocess.Popen", return_value=_make_proc()): + hc_main.hcatOllama("0", env.hash_file, "wordlist", str(wl)) + + captured = capsys.readouterr() + assert "Loaded 3" in captured.out + sample = gen.call_args[0][4]["sample"] + for line in sample.splitlines(): + assert line.strip() != "", "blank lines leaked into sample" + + +# --------------------------------------------------------------------------- +# Spinner is invoked (via the TTY guard — non-TTY in test suite) +# --------------------------------------------------------------------------- + + +def test_spinner_called_with_model_message(env, capsys): + """The spinner message must include the model name (non-TTY: just print).""" + wl = env.tmp_path / "s.txt" + wl.write_text("pw\n") + + with _ollama_globals(env.tmp_path), mock.patch.object( + hc_main.llm, "generate_candidates", return_value=["x"] + ), mock.patch("subprocess.Popen", return_value=_make_proc()): + hc_main.hcatOllama("0", env.hash_file, "wordlist", str(wl)) + + captured = capsys.readouterr() + assert MODEL in captured.out + assert "Generating password candidates via Ollama" in captured.out diff --git a/tests/test_progress_spinner.py b/tests/test_progress_spinner.py new file mode 100644 index 0000000..c59963c --- /dev/null +++ b/tests/test_progress_spinner.py @@ -0,0 +1,108 @@ +"""Tests for hate_crack/progress.py — the generic spinner context manager.""" + +from __future__ import annotations + +import io +import sys +import threading +import time +from unittest import mock + +import pytest + +from hate_crack.progress import spinner + + +# --------------------------------------------------------------------------- +# TTY guard: non-TTY path must print message once and not start a thread +# --------------------------------------------------------------------------- + + +def test_non_tty_prints_message(capsys: pytest.CaptureFixture[str]) -> None: + """In a non-TTY environment the message is printed once and no thread runs.""" + with mock.patch.object(sys, "stdout", wraps=sys.stdout) as m_stdout: + m_stdout.isatty.return_value = False # type: ignore[attr-defined] + # Use actual capsys for the print capture but override isatty + original_isatty = sys.stdout.isatty + sys.stdout.isatty = lambda: False # type: ignore[method-assign] + try: + with spinner("Working..."): + pass + finally: + sys.stdout.isatty = original_isatty # type: ignore[method-assign] + + captured = capsys.readouterr() + assert "Working..." in captured.out + + +def test_non_tty_no_extra_thread() -> None: + """Spinner in non-TTY context should not start a background thread.""" + buf = io.StringIO() + buf.isatty = lambda: False # type: ignore[attr-defined] + + threads_before = threading.active_count() + with mock.patch("sys.stdout", buf): + with spinner("No threads please"): + peak = threading.active_count() + # The thread count during the body should not exceed baseline + 1 + # (pytest itself may add threads; we just want no extra daemon spinner thread). + assert peak <= threads_before + 1 + + +# --------------------------------------------------------------------------- +# TTY path: spinner starts a thread, clears the line on exit +# --------------------------------------------------------------------------- + + +def test_tty_clears_line_on_exit() -> None: + """After the context exits the spinner line must be erased.""" + buf = io.StringIO() + buf.isatty = lambda: True # type: ignore[attr-defined] + + with mock.patch("sys.stdout", buf): + with spinner("Testing..."): + time.sleep(0.05) # let spinner tick at least once + + output = buf.getvalue() + # Line clear sequence must be present somewhere after the spinner started. + assert "\033[2K\r" in output or ("\033[2K" in output) + + +def test_tty_shows_elapsed_seconds() -> None: + """The spinner should show at least a 0s counter in its output.""" + buf = io.StringIO() + buf.isatty = lambda: True # type: ignore[attr-defined] + + with mock.patch("sys.stdout", buf): + with spinner("Counting..."): + time.sleep(0.25) # enough time for several ticks + + output = buf.getvalue() + # Should have written at least one elapsed counter like "0s" or "1s" + assert "s" in output + assert "Counting..." in output + + +def test_tty_exception_still_clears_line() -> None: + """Spinner must clear the line even when the body raises.""" + buf = io.StringIO() + buf.isatty = lambda: True # type: ignore[attr-defined] + + with mock.patch("sys.stdout", buf): + with pytest.raises(ValueError): + with spinner("Will raise"): + raise ValueError("boom") + + output = buf.getvalue() + assert "\033[2K\r" in output or ("\033[2K" in output) + + +def test_spinner_does_not_swallow_exception() -> None: + """The spinner context manager must re-raise exceptions from the body.""" + buf = io.StringIO() + buf.isatty = lambda: True # type: ignore[attr-defined] + + with mock.patch("sys.stdout", buf): + with pytest.raises(RuntimeError, match="test error"): + with spinner("Should propagate"): + raise RuntimeError("test error")