mirror of
https://github.com/trustedsec/hate_crack.git
synced 2026-07-28 14:47:22 -07:00
Justin BollingerandClaude
549c5a0a64
The Atomic Agents client was built with no timeout, so an Ollama server that accepted the TCP connection but never replied (most commonly a large model still loading into VRAM) left the CLI blocked in agent.run() forever. The caller's `except Exception` handler only ever fired on connection-refused. - llm.generate_candidates() takes a `timeout` parameter (default DEFAULT_TIMEOUT_SECONDS = 300.0) and forwards it to the OpenAI client. - openai.APITimeoutError is translated into a new domain-level llm.LLMTimeoutError so main.py need not import openai itself, keeping the atomic-agents/instructor dependency isolated to llm.py as documented. - hcatOllama passes the new `ollamaTimeout` config value and prints timeout-specific guidance (elapsed seconds, VRAM-loading hint, the setting to raise) instead of the misleading "ensure Ollama is running" message. - Documented `ollamaTimeout` in config.json.example and README. Also closes two test gaps: the defensive `except ValueError` handler in hcatOllama now has coverage, and the misleadingly-named `test_unknown_mode_raises` is split into explicit unit tests of _build_request's mode validation. Ticked the completed steps in the LLM Atomic Agents plan doc. Co-Authored-By: Claude <noreply@anthropic.com>
231 lines
9.4 KiB
Python
231 lines
9.4 KiB
Python
"""Orchestration tests for hcatOllama (candidate generation is mocked)."""
|
|
|
|
import os
|
|
from types import SimpleNamespace
|
|
from contextlib import contextmanager
|
|
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 = "qwen2.5:32b"
|
|
|
|
|
|
@pytest.fixture
|
|
def ollama_env(tmp_path):
|
|
hash_file = tmp_path / "hashes.txt"
|
|
hash_file.touch()
|
|
wordlist = tmp_path / "sample.txt"
|
|
wordlist.write_text("password\n123456\nletmein\n")
|
|
return SimpleNamespace(
|
|
tmp_path=tmp_path, hash_file=str(hash_file), wordlist=str(wordlist)
|
|
)
|
|
|
|
|
|
@contextmanager
|
|
def ollama_globals(tmp_path, tuning="", potfile=""):
|
|
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, "hcatBin", "/usr/bin/hashcat"), \
|
|
mock.patch.object(hc_main, "hcatTuning", tuning), \
|
|
mock.patch.object(hc_main, "hcatPotfilePath", potfile), \
|
|
mock.patch.object(hc_main, "rulesDirectory", rules_dir), \
|
|
mock.patch("hate_crack.main.generate_session_id", return_value="s"):
|
|
yield
|
|
|
|
|
|
def _make_proc(wait_return=0):
|
|
proc = mock.MagicMock()
|
|
proc.wait.return_value = wait_return
|
|
proc.communicate.return_value = (b"", b"")
|
|
proc.returncode = wait_return
|
|
return proc
|
|
|
|
|
|
def test_pull_ollama_model_is_gone():
|
|
assert not hasattr(hc_main, "_pull_ollama_model")
|
|
|
|
|
|
def test_target_mode_passes_dict_through(ollama_env):
|
|
with ollama_globals(ollama_env.tmp_path), \
|
|
mock.patch("hate_crack.main.llm.generate_candidates",
|
|
return_value=["Password1"]) as gen, \
|
|
mock.patch("subprocess.Popen", return_value=_make_proc()):
|
|
hc_main.hcatOllama(
|
|
"0", ollama_env.hash_file, "target",
|
|
{"company": "ACME", "industry": "tech", "location": "NYC"},
|
|
)
|
|
gen.assert_called_once()
|
|
args = gen.call_args[0]
|
|
assert args[0] == OLLAMA_URL and args[1] == MODEL and args[3] == "target"
|
|
assert args[4] == {"company": "ACME", "industry": "tech", "location": "NYC"}
|
|
|
|
|
|
def test_wordlist_mode_reads_file_and_passes_sample(ollama_env):
|
|
with ollama_globals(ollama_env.tmp_path), \
|
|
mock.patch("hate_crack.main.llm.generate_candidates",
|
|
return_value=["Password1"]) as gen, \
|
|
mock.patch("subprocess.Popen", return_value=_make_proc()):
|
|
hc_main.hcatOllama("0", ollama_env.hash_file, "wordlist", ollama_env.wordlist)
|
|
ctx_data = gen.call_args[0][4]
|
|
assert "letmein" in ctx_data["sample"]
|
|
|
|
|
|
def test_wordlist_mode_strips_hash_prefix(ollama_env):
|
|
# hash:password lines should contribute only the post-colon plaintext.
|
|
dump = ollama_env.tmp_path / "dump.txt"
|
|
dump.write_text("aad3b435:Winter2024\nnocolonline\n")
|
|
with ollama_globals(ollama_env.tmp_path), \
|
|
mock.patch("hate_crack.main.llm.generate_candidates",
|
|
return_value=["x"]) as gen, \
|
|
mock.patch("subprocess.Popen", return_value=_make_proc()):
|
|
hc_main.hcatOllama("0", ollama_env.hash_file, "wordlist", str(dump))
|
|
sample = gen.call_args[0][4]["sample"]
|
|
assert "Winter2024" in sample
|
|
assert "aad3b435" not in sample
|
|
assert "nocolonline" in sample
|
|
|
|
|
|
def test_per_rule_runs_hashcat_with_rule_flag(ollama_env):
|
|
rules_dir = str(ollama_env.tmp_path / "rules")
|
|
os.makedirs(rules_dir, exist_ok=True)
|
|
rule_file = os.path.join(rules_dir, "test.rule")
|
|
open(rule_file, "w").close()
|
|
calls = []
|
|
|
|
def track_popen(cmd, **kwargs):
|
|
calls.append(list(cmd))
|
|
return _make_proc()
|
|
|
|
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, "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"), \
|
|
mock.patch("hate_crack.main.llm.generate_candidates",
|
|
return_value=["Password1"]), \
|
|
mock.patch("subprocess.Popen", side_effect=track_popen):
|
|
hc_main.hcatOllama("1000", ollama_env.hash_file, "target",
|
|
{"company": "X", "industry": "Y", "location": "Z"})
|
|
|
|
# call 0 = plain wordlist run, call 1 = rule run
|
|
assert len(calls) == 2
|
|
rule_call = calls[1]
|
|
assert "-r" in rule_call
|
|
assert rule_file in rule_call
|
|
|
|
|
|
def test_missing_wordlist_prints_error(ollama_env, capsys):
|
|
with ollama_globals(ollama_env.tmp_path), \
|
|
mock.patch("hate_crack.main.llm.generate_candidates") as gen:
|
|
hc_main.hcatOllama("0", ollama_env.hash_file, "wordlist", "/no/such.txt")
|
|
captured = capsys.readouterr()
|
|
assert "Wordlist not found" in captured.out
|
|
gen.assert_not_called()
|
|
|
|
|
|
def test_writes_candidates_and_runs_hashcat(ollama_env):
|
|
calls = []
|
|
|
|
def track_popen(cmd, **kwargs):
|
|
calls.append(list(cmd))
|
|
return _make_proc()
|
|
|
|
with ollama_globals(ollama_env.tmp_path), \
|
|
mock.patch("hate_crack.main.llm.generate_candidates",
|
|
return_value=["Password1", "Summer2024"]), \
|
|
mock.patch("subprocess.Popen", side_effect=track_popen):
|
|
hc_main.hcatOllama("1000", ollama_env.hash_file, "target",
|
|
{"company": "X", "industry": "Y", "location": "Z"})
|
|
|
|
candidates_path = f"{ollama_env.hash_file}.ollama_candidates"
|
|
assert os.path.isfile(candidates_path)
|
|
with open(candidates_path) as f:
|
|
assert f.read().splitlines() == ["Password1", "Summer2024"]
|
|
# First hashcat call is the plain wordlist run with the candidates file.
|
|
assert calls and candidates_path in calls[0]
|
|
assert "-r" not in calls[0]
|
|
|
|
|
|
def test_empty_candidates_skips_hashcat(ollama_env, capsys):
|
|
with ollama_globals(ollama_env.tmp_path), \
|
|
mock.patch("hate_crack.main.llm.generate_candidates", return_value=[]), \
|
|
mock.patch("subprocess.Popen") as popen:
|
|
hc_main.hcatOllama("0", ollama_env.hash_file, "target",
|
|
{"company": "X", "industry": "Y", "location": "Z"})
|
|
captured = capsys.readouterr()
|
|
assert "no usable" in captured.out.lower()
|
|
popen.assert_not_called()
|
|
|
|
|
|
def test_generation_error_reports_and_aborts(ollama_env, capsys):
|
|
with ollama_globals(ollama_env.tmp_path), \
|
|
mock.patch("hate_crack.main.llm.generate_candidates",
|
|
side_effect=Exception("connection refused")), \
|
|
mock.patch("subprocess.Popen") as popen:
|
|
hc_main.hcatOllama("0", ollama_env.hash_file, "target",
|
|
{"company": "X", "industry": "Y", "location": "Z"})
|
|
captured = capsys.readouterr()
|
|
assert "Ensure Ollama is running" in captured.out
|
|
popen.assert_not_called()
|
|
|
|
|
|
def test_timeout_error_reports_timeout_guidance(ollama_env, capsys):
|
|
with ollama_globals(ollama_env.tmp_path), \
|
|
mock.patch.object(hc_main, "ollamaTimeout", 300.0), \
|
|
mock.patch("hate_crack.main.llm.generate_candidates",
|
|
side_effect=hc_main.llm.LLMTimeoutError("timed out")), \
|
|
mock.patch("subprocess.Popen") as popen:
|
|
hc_main.hcatOllama("0", ollama_env.hash_file, "target",
|
|
{"company": "X", "industry": "Y", "location": "Z"})
|
|
captured = capsys.readouterr()
|
|
assert "timed out" in captured.out.lower()
|
|
assert "300" in captured.out
|
|
assert "ollamaTimeout" in captured.out
|
|
assert "Ensure Ollama is running" not in captured.out
|
|
popen.assert_not_called()
|
|
assert not os.path.isfile(f"{ollama_env.hash_file}.ollama_candidates")
|
|
|
|
|
|
def test_value_error_reports_message_and_aborts(ollama_env, capsys):
|
|
"""Covers the defensive `except ValueError` handler in hcatOllama."""
|
|
with ollama_globals(ollama_env.tmp_path), \
|
|
mock.patch("hate_crack.main.llm.generate_candidates",
|
|
side_effect=ValueError("boom")), \
|
|
mock.patch("subprocess.Popen") as popen:
|
|
hc_main.hcatOllama("0", ollama_env.hash_file, "target",
|
|
{"company": "X", "industry": "Y", "location": "Z"})
|
|
captured = capsys.readouterr()
|
|
assert "Error: boom" in captured.out
|
|
popen.assert_not_called()
|
|
assert not os.path.isfile(f"{ollama_env.hash_file}.ollama_candidates")
|
|
|
|
|
|
def test_timeout_config_forwarded_to_generate_candidates(ollama_env):
|
|
with ollama_globals(ollama_env.tmp_path), \
|
|
mock.patch.object(hc_main, "ollamaTimeout", 77.0), \
|
|
mock.patch("hate_crack.main.llm.generate_candidates",
|
|
return_value=["Password1"]) as gen, \
|
|
mock.patch("subprocess.Popen", return_value=_make_proc()):
|
|
hc_main.hcatOllama("0", ollama_env.hash_file, "target",
|
|
{"company": "X", "industry": "Y", "location": "Z"})
|
|
assert gen.call_args.kwargs["timeout"] == 77.0
|
|
|
|
|
|
def test_unknown_mode_prints_error(ollama_env, capsys):
|
|
with ollama_globals(ollama_env.tmp_path), \
|
|
mock.patch("hate_crack.main.llm.generate_candidates") as gen:
|
|
hc_main.hcatOllama("0", ollama_env.hash_file, "bogus", {})
|
|
captured = capsys.readouterr()
|
|
assert "Unknown LLM generation mode" in captured.out
|
|
gen.assert_not_called()
|