Merge pull request #129 from trustedsec/feature/llm-atomic-agents

feat(llm): refactor LLM attack onto Atomic Agents, add wordlist mode (2.12.0)

CHANGELOG.md

@@ -7,6 +7,36 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 Dates are omitted for releases predating this file; see the git tags for exact timing.
 
+## [2.12.0] - 2026-07-24
+
+### Changed
+
+- **LLM attack now uses the Atomic Agents framework** for structured (JSON) candidate
+  generation instead of raw HTTP + regex line-parsing. Candidate generation lives in the
+  new `hate_crack/llm.py` module.
+- **Default Ollama model is now `qwen2.5:32b`** (was `mistral`), chosen for reliable
+  structured-output adherence.
+
+### Added
+
+- **Wordlist (denylist) generation mode** for the LLM attack is now reachable from the
+  menu: select the LLM attack (option 12), then choose "Wordlist" to derive basewords from
+  a sample wordlist.
+
+### Fixed
+
+- **The LLM attack no longer hangs forever waiting on Ollama.** Generation requests are now
+  bounded by a configurable timeout (`ollamaTimeout` in `config.json`, default 300 seconds).
+  Previously, if Ollama accepted the connection but never replied — most commonly a large
+  model still loading into VRAM — hate_crack sat at a frozen prompt with no recourse but
+  Ctrl-C. When the timeout fires you now get a specific message naming the elapsed timeout
+  and the setting to raise, instead of a misleading "ensure Ollama is running" hint.
+
+### Removed
+
+- **Automatic model pulling.** hate_crack no longer pulls missing Ollama models; pull them
+  yourself with `ollama pull <model>`.
+
 ## [2.11.4] - 2026-07-24
 
 ### Added

README.md

@@ -459,18 +459,22 @@ Set Hashview credentials in `config.json`:
 
 #### Ollama Configuration
 
-The LLM Attack (option 15) uses Ollama to generate password candidates. Configure the model and context window in `config.json`:
+The LLM Attack (option 12) uses Ollama to generate password candidates. Configure the model, context window, and request timeout in `config.json`:
 
 ```json
 {
-  "ollamaModel": "mistral",
-  "ollamaNumCtx": 2048
+  "ollamaModel": "qwen2.5:32b",
+  "ollamaNumCtx": 2048,
+  "ollamaTimeout": 300
 }
 ```
 
-- **`ollamaModel`** — The Ollama model to use for candidate generation (default: `mistral`).
+- **`ollamaModel`** — The Ollama model used for candidate generation (default: `qwen2.5:32b`). The LLM attack uses structured (JSON) output, so choose a model with good tool/JSON support.
 - **`ollamaNumCtx`** — Context window size for the model (default: `2048`).
-- The Ollama URL defaults to `http://localhost:11434`. Ensure Ollama is running before using the LLM Attack.
+- **`ollamaTimeout`** — Seconds to wait for a generation response before giving up (default: `300`). Raise this if a large model is still loading into VRAM on the first request, which can otherwise exceed the timeout; hate_crack prints the elapsed timeout and this setting's name when it fires.
+- The Ollama URL defaults to `http://localhost:11434` (override via the `OLLAMA_HOST` env var). Ensure Ollama is running and the model is pulled (`ollama pull qwen2.5:32b`) before using the LLM Attack — hate_crack no longer auto-pulls missing models.
+
+The attack offers two generation modes: **Target info** (company / industry / location) and **Wordlist** (derive denylist basewords from a sample wordlist).
 
 ### Notifications (menu option 82)
 

config.json.example

@@ -25,8 +25,9 @@
   "hashview_url": "http://localhost:8443",
   "hashview_api_key": "",
   "hashmob_api_key": "",
-  "ollamaModel": "mistral",
+  "ollamaModel": "qwen2.5:32b",
   "ollamaNumCtx": 2048,
+  "ollamaTimeout": 300,
   "omenTrainingList": "rockyou.txt",
   "omenMaxCandidates": 50000000,
   "pcfgRuleset": "DEFAULT",

hate_crack/attacks.py

@@ -513,33 +513,45 @@ def bandrel_method(ctx: Any) -> None:
 def ollama_attack(ctx: Any) -> None:
     _notify.prompt_notify_for_attack("LLM")
     print("\n\tLLM Attack")
-    company = input("Company name: ").strip()
-    industry = input("Industry: ").strip()
-    location = input("Location: ").strip()
-    target_info = {
-        "company": company,
-        "industry": industry,
-        "location": location,
-    }
-    ctx.hcatOllama(ctx.hcatHashType, ctx.hcatHashFile, "target", target_info)
+    print("\t1. Target info (company / industry / location)")
+    print("\t2. Wordlist (generate basewords from a sample wordlist)")
+    choice = input("\n\tSelect generation mode: ").strip()
+
+    if choice == "1":
+        company = input("Company name: ").strip()
+        industry = input("Industry: ").strip()
+        location = input("Location: ").strip()
+        ctx.hcatOllama(
+            ctx.hcatHashType,
+            ctx.hcatHashFile,
+            "target",
+            {"company": company, "industry": industry, "location": location},
+        )
+    elif choice == "2":
+        path = _omen_pick_training_wordlist(ctx, title="LLM Sample Wordlists")
+        if not path:
+            return
+        ctx.hcatOllama(ctx.hcatHashType, ctx.hcatHashFile, "wordlist", path)
+    else:
+        print("\t[!] Invalid selection.")
 
 
-def _omen_pick_training_wordlist(ctx: Any):
-    """Show wordlist picker for OMEN training. Returns path or None."""
+def _omen_pick_training_wordlist(ctx: Any, title: str = "Training Wordlists"):
+    """Show wordlist picker. Returns path or None."""
     wordlist_files = ctx.list_wordlist_files(ctx.hcatWordlists)
     if wordlist_files:
         entries = [f"{i}) {f}" for i, f in enumerate(wordlist_files, start=1)]
         max_len = max((len(e) for e in entries), default=24)
         print_multicolumn_list(
-            "Training Wordlists",
+            title,
             entries,
             min_col_width=max_len,
             max_col_width=max_len,
         )
     print("\tp. Enter a custom path")
-    sel = input("\n\tSelect wordlist for training: ").strip()
+    sel = input("\n\tSelect wordlist: ").strip()
     if sel.lower() == "p":
-        path = input("\n\tPath to training wordlist: ").strip()
+        path = input("\n\tPath to wordlist: ").strip()
         return path if path else None
     try:
         idx = int(sel)

hate_crack/llm.py

@@ -0,0 +1,152 @@
+"""Structured LLM password-candidate generation via Atomic Agents + Ollama.
+
+Isolates the atomic-agents / instructor dependency. The rest of hate_crack talks
+to this module only through ``generate_candidates``.
+"""
+
+import instructor
+from openai import APITimeoutError, OpenAI
+from pydantic import Field
+
+from atomic_agents import AgentConfig, AtomicAgent, BaseIOSchema
+from atomic_agents.context import SystemPromptGenerator
+
+MAX_CANDIDATE_LEN = 128
+DEFAULT_TIMEOUT_SECONDS = 300.0
+
+
+class LLMTimeoutError(Exception):
+    """The LLM server accepted the request but did not respond in time.
+
+    Raised instead of ``openai.APITimeoutError`` so callers do not need to import
+    ``openai`` themselves.
+    """
+
+
+class GenerationInput(BaseIOSchema):
+    """Instruction and context for a password-candidate generation request."""
+
+    request: str = Field(
+        ..., description="The full instruction, including any target or sample context."
+    )
+
+
+class PasswordCandidatesOutput(BaseIOSchema):
+    """A structured list of candidate passwords / basewords."""
+
+    candidates: list[str] = Field(
+        ...,
+        description=(
+            "Candidate passwords, one per list entry. No numbering, bullets, or "
+            "explanation — just the raw candidate strings."
+        ),
+    )
+
+
+_TARGET_PROMPT = SystemPromptGenerator(
+    background=[
+        "You are a security professional generating password candidates during an "
+        "authorized penetration test / capture-the-flag exercise.",
+    ],
+    steps=[
+        "Study the provided target context (company, industry, location).",
+        "Derive basewords from the company name and industry terms.",
+        "Combine basewords with common suffixes, years, and leetspeak substitutions.",
+    ],
+    output_instructions=[
+        "Return only candidate passwords in the candidates list.",
+        "Do not include explanations, numbering, or duplicate entries.",
+    ],
+)
+
+_WORDLIST_PROMPT = SystemPromptGenerator(
+    background=[
+        "You build denylist basewords so users cannot set weak passwords.",
+    ],
+    steps=[
+        "Study the sample passwords for patterns: capitalization, leetspeak, "
+        "suffixes, and common substitutions.",
+        "Produce basewords that capture those patterns.",
+    ],
+    output_instructions=[
+        "Return only basewords in the candidates list.",
+        "Do not include explanations, numbering, or duplicate entries.",
+    ],
+)
+
+
+def _build_request(mode: str, context_data: dict) -> str:
+    """Build the natural-language request string for the given mode."""
+    if mode == "target":
+        company = context_data.get("company", "")
+        industry = context_data.get("industry", "")
+        location = context_data.get("location", "")
+        return (
+            f"The target organization is '{company}', a {industry} in {location}. "
+            "Generate as many plausible password candidates as you can, using "
+            "permutations of the company name and industry terms with common "
+            "suffixes, years, and leetspeak substitutions."
+        )
+    if mode == "wordlist":
+        sample = context_data.get("sample", "")
+        return (
+            "Here are sample passwords. Study their patterns and generate basewords "
+            "for a denylist:\n" + sample
+        )
+    raise ValueError(f"Unknown LLM generation mode: {mode}")
+
+
+def generate_candidates(
+    url: str,
+    model: str,
+    num_ctx: int,
+    mode: str,
+    context_data: dict,
+    timeout: float = DEFAULT_TIMEOUT_SECONDS,
+) -> list[str]:
+    """Generate password candidates via an Ollama-backed AtomicAgent.
+
+    ``timeout`` is the number of seconds to wait for a generation response before
+    giving up; it bounds the whole request so a server that accepts the
+    connection but never replies (e.g. a large model still loading into VRAM)
+    cannot hang the caller forever.
+
+    Returns a deduped, length-capped list of candidate strings (may be empty).
+    Raises ValueError for an unknown mode and LLMTimeoutError if the request
+    exceeds ``timeout``. Other client/connection errors propagate to the caller.
+    """
+    request = _build_request(mode, context_data)
+
+    client = instructor.from_openai(
+        OpenAI(base_url=f"{url}/v1", api_key="ollama", timeout=timeout),
+        mode=instructor.Mode.JSON,
+    )
+    prompt_generator = _TARGET_PROMPT if mode == "target" else _WORDLIST_PROMPT
+
+    agent = AtomicAgent[GenerationInput, PasswordCandidatesOutput](
+        config=AgentConfig(
+            client=client,
+            model=model,
+            system_prompt_generator=prompt_generator,
+            model_api_parameters={"extra_body": {"options": {"num_ctx": num_ctx}}},
+        )
+    )
+
+    try:
+        result = agent.run(GenerationInput(request=request))
+    except APITimeoutError as e:
+        raise LLMTimeoutError(
+            f"no response from {url} within {timeout:g} seconds"
+        ) from e
+
+    seen: set[str] = set()
+    candidates: list[str] = []
+    for raw in getattr(result, "candidates", []) or []:
+        candidate = str(raw).strip()
+        if not candidate or len(candidate) > MAX_CANDIDATE_LEN:
+            continue
+        if candidate in seen:
+            continue
+        seen.add(candidate)
+        candidates.append(candidate)
+    return candidates

hate_crack/main.py

@@ -21,8 +21,6 @@ import subprocess
 import shlex
 import time
 import argparse
-import urllib.request
-import urllib.error
 import contextlib
 import gzip
 import lzma
@@ -75,6 +73,7 @@ from hate_crack.cli import (  # noqa: E402
     setup_logging,
 )
 from hate_crack import attacks as _attacks  # noqa: E402
+from hate_crack import llm  # noqa: E402
 from hate_crack.menu import interactive_menu  # noqa: E402
 from hate_crack.username_detect import detect_username_hash_format  # noqa: E402
 
@@ -451,8 +450,9 @@ hcatGoodMeasureBaseList = config_parser["hcatGoodMeasureBaseList"]
 hcatDebugLogPath = os.path.expanduser(config_parser["hcatDebugLogPath"])
 
 ollamaUrl = "http://" + os.environ.get("OLLAMA_HOST", "localhost:11434")
-ollamaModel = config_parser.get("ollamaModel", "mistral")
+ollamaModel = config_parser.get("ollamaModel", "qwen2.5:32b")
 ollamaNumCtx = int(config_parser.get("ollamaNumCtx", 2048))
+ollamaTimeout = float(config_parser.get("ollamaTimeout", 300))
 
 omenTrainingList = config_parser.get("omenTrainingList", "rockyou.txt")
 omenMaxCandidates = int(config_parser.get("omenMaxCandidates", 1000000))
@@ -2036,49 +2036,11 @@ def hcatBandrel(hcatHashType, hcatHashFile):
         _run_hcat_cmd(cmd, attack_name="Bandrel", hash_file=hcatHashFile)
 
 
-# Pull an Ollama model via the /api/pull streaming endpoint
-def _pull_ollama_model(url, model):
-    """Pull an Ollama model. Returns True on success, False on failure."""
-    print(f"Model '{model}' not found locally. Pulling from Ollama...")
-    pull_url = f"{url}/api/pull"
-    payload = json.dumps({"name": model, "stream": True}).encode("utf-8")
-    req = urllib.request.Request(
-        pull_url,
-        data=payload,
-        headers={"Content-Type": "application/json"},
-    )
-    try:
-        with urllib.request.urlopen(req) as resp:
-            for raw_line in resp:
-                line = raw_line.decode("utf-8", errors="replace").strip()
-                if not line:
-                    continue
-                try:
-                    data = json.loads(line)
-                except json.JSONDecodeError:
-                    continue
-                status = data.get("status")
-                if status:
-                    print(f"  {status}")
-    except urllib.error.HTTPError as e:
-        print(f"Error pulling model: HTTP {e.code}")
-        return False
-    except urllib.error.URLError as e:
-        print(f"Error: Could not connect to Ollama: {e}")
-        return False
-    except Exception as e:
-        print(f"Error pulling model: {e}")
-        return False
-    print(f"Successfully pulled model '{model}'.")
-    return True
-
-
 # LLM Ollama Attack
 def hcatOllama(hcatHashType, hcatHashFile, mode, context_data):
-    global hcatProcess
     candidates_path = f"{hcatHashFile}.ollama_candidates"
 
-    # Step A: Build LLM prompt based on mode
+    # Step A: normalize context into the dict generate_candidates expects.
     if mode == "wordlist":
         wordlist_path = context_data
         if not os.path.isfile(wordlist_path):
@@ -2091,7 +2053,7 @@ def hcatOllama(hcatHashType, hcatHashFile, mode, context_data):
                     stripped = line.strip()
                     if not stripped:
                         continue
-                    # Use only content after the first colon (e.g. hash:password -> password)
+                    # hash:password -> password
                     if ":" in stripped:
                         stripped = stripped.split(":", 1)[1]
                     if stripped:
@@ -2100,107 +2062,43 @@ def hcatOllama(hcatHashType, hcatHashFile, mode, context_data):
             print(f"Error reading wordlist: {e}")
             return
         print(f"Loaded {len(lines)} passwords from wordlist.")
-        wordlist_sample = "\n".join(lines)
-        prompt = (
-            "Generate baseword to be used in a denylist for keeping users from setting their passwords with these basewords."
-            "Study the patterns, character choices, and structures. Focus on patterns like capitalization, leetspeak, suffixes, and common substitutions. Here are the sample passwords:\n"
-            f"{wordlist_sample}"
-        )
+        gen_context = {"sample": "\n".join(lines)}
     elif mode == "target":
-        company = context_data.get("company", "")
-        industry = context_data.get("industry", "")
-        location = context_data.get("location", "")
-        prompt = (
-            "You are participating in a capture the flag event as a security professional. "
-            "You are my partner in the competition. You need to recover the password to a system to retrieve the flag. "
-            "Output as many possible password combinations you think might help us. "
-            f"The name of the fake company is {company}. They are a {industry} in {location}. "
-            "Use terms related to the industry as basewords and also use permutations of the company name combined with common suffixes. "
-            "Only output the candidate password each on a new line. Dont output any explanation. "
-            "Only output the password candidate. Do not number the lines or add any extra information to the output"
-        )
+        gen_context = context_data
     else:
         print(f"Error: Unknown LLM generation mode: {mode}")
         return
 
-    # Step B: Call Ollama API to generate candidates
+    # Step B: generate candidates via the Atomic Agents module.
     print(f"Generating password candidates via Ollama ({ollamaModel})...")
-    api_url = f"{ollamaUrl}/api/generate"
-    payload = json.dumps(
-        {
-            "model": ollamaModel,
-            "prompt": prompt,
-            "stream": False,
-            "options": {"num_ctx": ollamaNumCtx},
-        }
-    ).encode("utf-8")
-
-    if debug_mode:
-        print(f"[DEBUG] Ollama API URL: {api_url}")
-        print(f"[DEBUG] Ollama request payload: {payload.decode('utf-8')}")
-
     try:
-        req = urllib.request.Request(
-            api_url,
-            data=payload,
-            headers={"Content-Type": "application/json"},
+        candidates = llm.generate_candidates(
+            ollamaUrl,
+            ollamaModel,
+            ollamaNumCtx,
+            mode,
+            gen_context,
+            timeout=ollamaTimeout,
         )
-        with urllib.request.urlopen(req, timeout=600) as resp:
-            result = json.loads(resp.read().decode("utf-8"))
-        if debug_mode:
-            print(f"[DEBUG] Ollama response: {json.dumps(result, indent=2)}")
-    except urllib.error.HTTPError as e:
-        if e.code == 404:
-            if _pull_ollama_model(ollamaUrl, ollamaModel):
-                try:
-                    req = urllib.request.Request(
-                        api_url,
-                        data=payload,
-                        headers={"Content-Type": "application/json"},
-                    )
-                    with urllib.request.urlopen(req, timeout=600) as resp:
-                        result = json.loads(resp.read().decode("utf-8"))
-                    if debug_mode:
-                        print(
-                            f"[DEBUG] Ollama response (after pull): {json.dumps(result, indent=2)}"
-                        )
-                except Exception as retry_err:
-                    print(f"Error calling Ollama API after pull: {retry_err}")
-                    return
-            else:
-                print(f"Could not pull model '{ollamaModel}'. Aborting LLM attack.")
-                return
-        else:
-            print(f"Error: Could not connect to Ollama at {ollamaUrl}: {e}")
-            print("Ensure Ollama is running (ollama serve) and try again.")
-            return
-    except urllib.error.URLError as e:
-        print(f"Error: Could not connect to Ollama at {ollamaUrl}: {e}")
-        print("Ensure Ollama is running (ollama serve) and try again.")
+    except llm.LLMTimeoutError:
+        print(f"Error: the Ollama request timed out after {ollamaTimeout:g} seconds.")
+        print(
+            f"The model ({ollamaModel}) may still be loading into VRAM. Retry, or "
+            'raise "ollamaTimeout" in config.json to wait longer.'
+        )
+        return
+    except ValueError as e:
+        # Defensive: mode is already validated above, but keep an explicit,
+        # non-misleading message if generate_candidates ever rejects its input.
+        print(f"Error: {e}")
         return
     except Exception as e:
-        print(f"Error calling Ollama API: {e}")
-        return
-
-    response_text = result.get("response", "")
-    if "I'm sorry, but I can't help with that" in response_text:
+        print(f"Error generating candidates: {e}")
         print(
-            "Error: Ollama refused the request. Try a different model or adjust your prompt."
+            "Ensure Ollama is running (ollama serve) and the model is pulled "
+            f"(ollama pull {ollamaModel})."
         )
         return
-    raw_lines = response_text.strip().split("\n")
-    # Filter out blank lines and lines that look like numbering/explanation
-    candidates = []
-    for line in raw_lines:
-        stripped = line.strip()
-        if not stripped:
-            continue
-        # Strip leading numbering like "1. " or "1) " or "- "
-        cleaned = re.sub(r"^\d+[.)]\s*", "", stripped)
-        cleaned = re.sub(r"^[-*]\s*", "", cleaned)
-        cleaned = cleaned.strip()
-        if cleaned and len(cleaned) <= 128:
-            candidates.append(cleaned)
 
     if not candidates:
         print("Error: Ollama returned no usable password candidates.")
@@ -2215,13 +2113,8 @@ def hcatOllama(hcatHashType, hcatHashFile, mode, context_data):
         return
 
     print(f"Generated {len(candidates)} password candidates -> {candidates_path}")
-    if debug_mode:
-        filtered_count = len(raw_lines) - len(candidates)
-        print(
-            f"[DEBUG] Filtered out {filtered_count} lines from Ollama response ({len(raw_lines)} raw -> {len(candidates)} candidates)"
-        )
 
-    # Step C: Run hashcat wordlist attack with LLM-generated candidates (no rules)
+    # Step C: hashcat wordlist attack with the generated candidates (no rules).
     print("Running wordlist attack with LLM-generated candidates...")
     cmd = [
         hcatBin,
@@ -2246,7 +2139,7 @@ def hcatOllama(hcatHashType, hcatHashFile, mode, context_data):
     except KeyboardInterrupt:
         return
 
-    # Step D: Run hashcat with LLM candidates against every rule in the rules directory
+    # Step D: hashcat with candidates against every rule in the rules directory.
     rule_files = sorted(f for f in os.listdir(rulesDirectory) if f != ".DS_Store")
     if not rule_files:
         print("No rule files found in rules directory. Skipping rule-based attacks.")

pyproject.toml

@@ -15,6 +15,12 @@ dependencies = [
     "packaging>=26.2",
     "simple-term-menu==1.6.6",
     "click>=8.4.2",
+    "atomic-agents>=2.0.0",
+    # Imported directly by hate_crack/llm.py; declared explicitly rather than
+    # relying on atomic-agents pulling them in transitively.
+    "instructor>=1.14.5",
+    "openai>=2.48.0",
+    "pydantic>=2.13.4",
 ]
 
 [project.scripts]

tests/test_attacks_behavior.py

@@ -329,7 +329,7 @@ class TestOllamaAttack:
     def test_calls_hcatOllama_with_context(self) -> None:
         ctx = _make_ctx()
 
-        with patch("builtins.input", side_effect=["ACME", "tech", "NYC"]):
+        with patch("builtins.input", side_effect=["1", "ACME", "tech", "NYC"]):
             ollama_attack(ctx)
 
         ctx.hcatOllama.assert_called_once_with(
@@ -342,7 +342,7 @@ class TestOllamaAttack:
     def test_passes_hash_type_and_file(self) -> None:
         ctx = _make_ctx(hash_type="1800", hash_file="/tmp/sha512.txt")
 
-        with patch("builtins.input", side_effect=["Corp", "finance", "London"]):
+        with patch("builtins.input", side_effect=["1", "Corp", "finance", "London"]):
             ollama_attack(ctx)
 
         call_args = ctx.hcatOllama.call_args[0]
@@ -352,7 +352,7 @@ class TestOllamaAttack:
     def test_strips_whitespace_from_inputs(self) -> None:
         ctx = _make_ctx()
 
-        with patch("builtins.input", side_effect=["  ACME  ", "  tech  ", "  NYC  "]):
+        with patch("builtins.input", side_effect=["1", "  ACME  ", "  tech  ", "  NYC  "]):
             ollama_attack(ctx)
 
         target_info = ctx.hcatOllama.call_args[0][3]
@@ -363,7 +363,34 @@ class TestOllamaAttack:
     def test_target_string_is_literal_target(self) -> None:
         ctx = _make_ctx()
 
-        with patch("builtins.input", side_effect=["X", "Y", "Z"]):
+        with patch("builtins.input", side_effect=["1", "X", "Y", "Z"]):
             ollama_attack(ctx)
 
         assert ctx.hcatOllama.call_args[0][2] == "target"
+
+    def test_wordlist_mode_calls_hcatOllama_with_path(self) -> None:
+        ctx = _make_ctx()
+        ctx.list_wordlist_files.return_value = ["rockyou.txt"]
+        ctx.hcatWordlists = "/tmp/wl"
+
+        # mode "2", then pick wordlist "1"
+        with patch("builtins.input", side_effect=["2", "1"]):
+            ollama_attack(ctx)
+
+        args = ctx.hcatOllama.call_args[0]
+        assert args[2] == "wordlist"
+        assert args[3].endswith("rockyou.txt")
+
+    def test_invalid_mode_does_not_call_hcatOllama(self) -> None:
+        ctx = _make_ctx()
+        with patch("builtins.input", side_effect=["9"]):
+            ollama_attack(ctx)
+        ctx.hcatOllama.assert_not_called()
+
+    def test_wordlist_mode_aborts_when_no_wordlist_picked(self) -> None:
+        ctx = _make_ctx()
+        # mode "2", then an invalid picker selection -> picker returns None
+        ctx.list_wordlist_files.return_value = []
+        with patch("builtins.input", side_effect=["2", "nonsense"]):
+            ollama_attack(ctx)
+        ctx.hcatOllama.assert_not_called()

tests/test_hcat_ollama.py

@@ -0,0 +1,230 @@
+"""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()

tests/test_llm.py

@@ -0,0 +1,147 @@
+"""Unit tests for hate_crack.llm.generate_candidates."""
+
+import os
+from unittest import mock
+
+import httpx
+import instructor
+import openai
+import pytest
+
+os.environ["HATE_CRACK_SKIP_INIT"] = "1"
+from hate_crack import llm  # noqa: E402
+
+
+def _patch_agent(candidates):
+    """Patch the client builders + AtomicAgent so no network happens.
+
+    Returns the AtomicAgent class mock so callers can inspect construction.
+    """
+    result = mock.MagicMock()
+    result.candidates = list(candidates)
+
+    agent_instance = mock.MagicMock()
+    agent_instance.run.return_value = result
+
+    agent_cls = mock.MagicMock()
+    # AtomicAgent[In, Out](config=...) -> agent_instance
+    agent_cls.__getitem__.return_value.return_value = agent_instance
+
+    return (
+        mock.patch("hate_crack.llm.instructor.from_openai", return_value=mock.MagicMock(spec=instructor.Instructor)),
+        mock.patch("hate_crack.llm.OpenAI"),
+        mock.patch("hate_crack.llm.AtomicAgent", agent_cls),
+        agent_cls,
+        agent_instance,
+    )
+
+
+def test_target_mode_returns_candidates():
+    p_instr, p_openai, p_agent, agent_cls, agent_instance = _patch_agent(
+        ["AcmeCorp2024", "Finance123"]
+    )
+    with p_instr, p_openai, p_agent:
+        out = llm.generate_candidates(
+            "http://localhost:11434",
+            "qwen2.5:32b",
+            2048,
+            "target",
+            {"company": "AcmeCorp", "industry": "Finance", "location": "NYC"},
+        )
+    assert out == ["AcmeCorp2024", "Finance123"]
+    # The instruction the agent received must include the target context.
+    run_arg = agent_instance.run.call_args[0][0]
+    assert "AcmeCorp" in run_arg.request
+    assert "Finance" in run_arg.request
+    assert "NYC" in run_arg.request
+
+
+def test_wordlist_mode_includes_sample_in_request():
+    p_instr, p_openai, p_agent, agent_cls, agent_instance = _patch_agent(["Passw0rd"])
+    with p_instr, p_openai, p_agent:
+        llm.generate_candidates(
+            "http://localhost:11434",
+            "qwen2.5:32b",
+            2048,
+            "wordlist",
+            {"sample": "password\nletmein\nsummer2024"},
+        )
+    run_arg = agent_instance.run.call_args[0][0]
+    assert "letmein" in run_arg.request
+
+
+def test_dedupes_and_caps_length():
+    long_pw = "A" * 129
+    p_instr, p_openai, p_agent, agent_cls, agent_instance = _patch_agent(
+        ["  keep  ", "keep", "dup", "dup", long_pw, ""]
+    )
+    with p_instr, p_openai, p_agent:
+        out = llm.generate_candidates(
+            "http://localhost:11434", "qwen2.5:32b", 2048,
+            "target", {"company": "X", "industry": "Y", "location": "Z"},
+        )
+    assert out == ["keep", "dup"]  # trimmed, deduped, blank + >128 dropped
+
+
+def test_num_ctx_forwarded_via_model_api_parameters():
+    p_instr, p_openai, p_agent, agent_cls, agent_instance = _patch_agent(["x"])
+    with p_instr, p_openai, p_agent:
+        llm.generate_candidates(
+            "http://localhost:11434", "qwen2.5:32b", 4096,
+            "target", {"company": "X", "industry": "Y", "location": "Z"},
+        )
+    # AtomicAgent[In, Out](config=<AgentConfig>) — inspect the config.
+    config = agent_cls.__getitem__.return_value.call_args.kwargs["config"]
+    assert config.model == "qwen2.5:32b"
+    assert config.model_api_parameters["extra_body"]["options"]["num_ctx"] == 4096
+
+
+def test_build_request_rejects_unknown_mode():
+    """Unit test of _build_request's mode validation only — no agent involved."""
+    with pytest.raises(ValueError, match="Unknown LLM generation mode: bogus"):
+        llm._build_request("bogus", {})
+
+
+def test_generate_candidates_rejects_unknown_mode_before_building_client():
+    """generate_candidates surfaces _build_request's ValueError to its caller."""
+    with pytest.raises(ValueError, match="Unknown LLM generation mode: bogus"):
+        llm.generate_candidates(
+            "http://localhost:11434", "qwen2.5:32b", 2048, "bogus", {},
+        )
+
+
+def test_timeout_forwarded_to_openai_client():
+    p_instr, p_openai, p_agent, agent_cls, agent_instance = _patch_agent(["x"])
+    with p_instr, p_openai as openai_cls, p_agent:
+        llm.generate_candidates(
+            "http://localhost:11434", "qwen2.5:32b", 2048,
+            "target", {"company": "X", "industry": "Y", "location": "Z"},
+            timeout=42.5,
+        )
+    assert openai_cls.call_args.kwargs["timeout"] == 42.5
+
+
+def test_default_timeout_used_when_omitted():
+    p_instr, p_openai, p_agent, agent_cls, agent_instance = _patch_agent(["x"])
+    with p_instr, p_openai as openai_cls, p_agent:
+        llm.generate_candidates(
+            "http://localhost:11434", "qwen2.5:32b", 2048,
+            "target", {"company": "X", "industry": "Y", "location": "Z"},
+        )
+    assert llm.DEFAULT_TIMEOUT_SECONDS == 300.0
+    assert openai_cls.call_args.kwargs["timeout"] == llm.DEFAULT_TIMEOUT_SECONDS
+
+
+def test_api_timeout_reraised_as_domain_error():
+    """openai.APITimeoutError is translated into llm.LLMTimeoutError."""
+    p_instr, p_openai, p_agent, agent_cls, agent_instance = _patch_agent(["x"])
+    agent_instance.run.side_effect = openai.APITimeoutError(
+        request=httpx.Request("POST", "http://localhost:11434/v1/chat/completions")
+    )
+    with p_instr, p_openai, p_agent:
+        with pytest.raises(llm.LLMTimeoutError):
+            llm.generate_candidates(
+                "http://localhost:11434", "qwen2.5:32b", 2048,
+                "target", {"company": "X", "industry": "Y", "location": "Z"},
+                timeout=1.0,
+            )

tests/test_pull_ollama_model.py

@@ -1,684 +0,0 @@
-"""Unit tests for _pull_ollama_model helper, hcatOllama steps A-C, and ollama_attack handler."""
-
-import io
-import json
-import os
-import urllib.error
-import urllib.request
-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
-
-
-# ---------------------------------------------------------------------------
-# Shared test infrastructure
-# ---------------------------------------------------------------------------
-
-OLLAMA_URL = "http://localhost:11434"
-MODEL = "llama3.2"
-
-
-@pytest.fixture
-def ollama_env(tmp_path):
-    """Create the filesystem layout hcatOllama expects."""
-    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=""):
-    """Patch the hc_main globals that hcatOllama reads."""
-    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, "hcatBin", "/usr/bin/hashcat"), \
-         mock.patch.object(hc_main, "hcatTuning", tuning), \
-         mock.patch.object(hc_main, "hcatPotfilePath", potfile), \
-         mock.patch.object(hc_main, "hate_path", str(tmp_path)), \
-         mock.patch.object(hc_main, "rulesDirectory", rules_dir), \
-         mock.patch("hate_crack.main.generate_session_id", return_value="test_session"):
-        yield
-
-
-def _make_proc(wait_return=0):
-    """Create a mock subprocess that works with both wait() and communicate()."""
-    proc = mock.MagicMock()
-    proc.wait.return_value = wait_return
-    proc.communicate.return_value = (b"", b"")
-    proc.returncode = wait_return
-    return proc
-
-
-def _generate_response(passwords):
-    """Build a fake urlopen context-manager that returns an Ollama /api/generate JSON body."""
-    body = json.dumps({"response": "\n".join(passwords)}).encode()
-    resp = mock.MagicMock()
-    resp.__enter__ = mock.Mock(return_value=io.BytesIO(body))
-    resp.__exit__ = mock.Mock(return_value=False)
-    return resp
-
-
-def _urlopen_with_response(passwords):
-    """Return a urlopen mock that always succeeds with the given passwords."""
-    return mock.patch(
-        "hate_crack.main.urllib.request.urlopen",
-        return_value=_generate_response(passwords),
-    )
-
-
-# ---------------------------------------------------------------------------
-# _pull_ollama_model tests
-# ---------------------------------------------------------------------------
-
-class TestPullOllamaModel:
-    """Tests for _pull_ollama_model()."""
-
-    OLLAMA_URL = "http://localhost:11434"
-    MODEL = "llama3.2"
-
-    def _make_stream_response(self, statuses):
-        """Build a fake streaming response (newline-delimited JSON)."""
-        lines = [json.dumps({"status": s}).encode() + b"\n" for s in statuses]
-        return io.BytesIO(b"".join(lines))
-
-    @mock.patch("hate_crack.main.urllib.request.urlopen")
-    def test_successful_pull(self, mock_urlopen, capsys):
-        mock_urlopen.return_value.__enter__ = mock.Mock(
-            return_value=self._make_stream_response(
-                ["pulling manifest", "downloading sha256:abc123", "success"]
-            )
-        )
-        mock_urlopen.return_value.__exit__ = mock.Mock(return_value=False)
-
-        result = hc_main._pull_ollama_model(self.OLLAMA_URL, self.MODEL)
-
-        assert result is True
-        captured = capsys.readouterr()
-        assert "not found locally" in captured.out
-        assert "Successfully pulled" in captured.out
-        assert "pulling manifest" in captured.out
-
-    @mock.patch("hate_crack.main.urllib.request.urlopen")
-    def test_pull_http_error(self, mock_urlopen, capsys):
-        mock_urlopen.side_effect = urllib.error.HTTPError(
-            url="http://localhost:11434/api/pull",
-            code=500,
-            msg="Internal Server Error",
-            hdrs=None,
-            fp=None,
-        )
-
-        result = hc_main._pull_ollama_model(self.OLLAMA_URL, self.MODEL)
-
-        assert result is False
-        captured = capsys.readouterr()
-        assert "HTTP 500" in captured.out
-
-    @mock.patch("hate_crack.main.urllib.request.urlopen")
-    def test_pull_url_error(self, mock_urlopen, capsys):
-        mock_urlopen.side_effect = urllib.error.URLError("Connection refused")
-
-        result = hc_main._pull_ollama_model(self.OLLAMA_URL, self.MODEL)
-
-        assert result is False
-        captured = capsys.readouterr()
-        assert "Could not connect" in captured.out
-
-    @mock.patch("hate_crack.main.urllib.request.urlopen")
-    def test_pull_generic_exception(self, mock_urlopen, capsys):
-        mock_urlopen.side_effect = RuntimeError("unexpected")
-
-        result = hc_main._pull_ollama_model(self.OLLAMA_URL, self.MODEL)
-
-        assert result is False
-        captured = capsys.readouterr()
-        assert "unexpected" in captured.out
-
-    @mock.patch("hate_crack.main.urllib.request.urlopen")
-    def test_pull_handles_empty_status_lines(self, mock_urlopen, capsys):
-        """Blank lines and missing status keys should not crash."""
-        lines = b'{"status": "downloading"}\n\n{"other": "data"}\n'
-        mock_urlopen.return_value.__enter__ = mock.Mock(
-            return_value=io.BytesIO(lines)
-        )
-        mock_urlopen.return_value.__exit__ = mock.Mock(return_value=False)
-
-        result = hc_main._pull_ollama_model(self.OLLAMA_URL, self.MODEL)
-
-        assert result is True
-
-    @mock.patch("hate_crack.main.urllib.request.urlopen")
-    def test_pull_request_payload(self, mock_urlopen):
-        """Verify the pull request sends the correct model name and stream=True."""
-        mock_urlopen.return_value.__enter__ = mock.Mock(
-            return_value=self._make_stream_response(["success"])
-        )
-        mock_urlopen.return_value.__exit__ = mock.Mock(return_value=False)
-
-        hc_main._pull_ollama_model(self.OLLAMA_URL, self.MODEL)
-
-        call_args = mock_urlopen.call_args
-        req = call_args[0][0]
-        body = json.loads(req.data.decode("utf-8"))
-        assert body["name"] == self.MODEL
-        assert body["stream"] is True
-        assert req.full_url == f"{self.OLLAMA_URL}/api/pull"
-
-
-# ---------------------------------------------------------------------------
-# hcatOllama 404-retry integration tests
-# ---------------------------------------------------------------------------
-
-class TestHcatOllama404Retry:
-    """Test that hcatOllama auto-pulls on 404 and retries the generate call."""
-
-    OLLAMA_URL = "http://localhost:11434"
-    MODEL = "llama3.2"
-
-    def _generate_response(self, passwords):
-        body = json.dumps({"response": "\n".join(passwords)}).encode()
-        return io.BytesIO(body)
-
-    def _setup_env(self, tmp_path):
-        """Create the files/dirs hcatOllama needs before it hits the API."""
-        hash_file = str(tmp_path / "hashes.txt")
-        open(hash_file, "w").close()
-
-        # Create a small wordlist for "wordlist" mode
-        wordlist = str(tmp_path / "sample.txt")
-        with open(wordlist, "w") as f:
-            f.write("password\n123456\nletmein\n")
-
-        return hash_file, wordlist
-
-    @mock.patch("hate_crack.main._pull_ollama_model")
-    @mock.patch("hate_crack.main.urllib.request.urlopen")
-    def test_404_triggers_pull_then_retries(self, mock_urlopen, mock_pull, capsys, tmp_path):
-        """A 404 on generate should trigger a pull, then retry successfully."""
-        hash_file, wordlist = self._setup_env(tmp_path)
-
-        # First call: 404, second call: success
-        generate_ok = mock.MagicMock()
-        generate_ok.__enter__ = mock.Mock(
-            return_value=self._generate_response(["Password1", "Summer2024"])
-        )
-        generate_ok.__exit__ = mock.Mock(return_value=False)
-
-        mock_urlopen.side_effect = [
-            urllib.error.HTTPError(
-                url=f"{self.OLLAMA_URL}/api/generate",
-                code=404,
-                msg="Not Found",
-                hdrs=None,
-                fp=None,
-            ),
-            generate_ok,
-        ]
-        mock_pull.return_value = True
-
-        rules_dir = str(tmp_path / "rules")
-        os.makedirs(rules_dir, exist_ok=True)
-        with mock.patch.object(hc_main, "ollamaUrl", self.OLLAMA_URL), \
-             mock.patch.object(hc_main, "ollamaModel", self.MODEL), \
-             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, "hate_path", str(tmp_path)), \
-             mock.patch.object(hc_main, "rulesDirectory", rules_dir), \
-             mock.patch("hate_crack.main.generate_session_id", return_value="test_session"), \
-             mock.patch("subprocess.Popen") as mock_popen:
-
-            proc = _make_proc()
-            mock_popen.return_value = proc
-
-            hc_main.hcatOllama("0", hash_file, "wordlist", wordlist)
-
-        mock_pull.assert_called_once_with(self.OLLAMA_URL, self.MODEL)
-        # urlopen called twice: first 404, then retry
-        assert mock_urlopen.call_count == 2
-
-    @mock.patch("hate_crack.main._pull_ollama_model")
-    @mock.patch("hate_crack.main.urllib.request.urlopen")
-    def test_404_pull_fails_aborts(self, mock_urlopen, mock_pull, capsys, tmp_path):
-        """If pull fails after 404, hcatOllama should abort gracefully."""
-        hash_file, wordlist = self._setup_env(tmp_path)
-
-        mock_urlopen.side_effect = urllib.error.HTTPError(
-            url=f"{self.OLLAMA_URL}/api/generate",
-            code=404,
-            msg="Not Found",
-            hdrs=None,
-            fp=None,
-        )
-        mock_pull.return_value = False
-
-        with mock.patch.object(hc_main, "ollamaUrl", self.OLLAMA_URL), \
-             mock.patch.object(hc_main, "ollamaModel", self.MODEL), \
-             mock.patch.object(hc_main, "hate_path", str(tmp_path)):
-
-            hc_main.hcatOllama("0", hash_file, "wordlist", wordlist)
-
-        captured = capsys.readouterr()
-        assert "Could not pull model" in captured.out
-        mock_pull.assert_called_once()
-
-    @mock.patch("hate_crack.main.urllib.request.urlopen")
-    def test_non_404_http_error_propagates(self, mock_urlopen, capsys, tmp_path):
-        """Non-404 HTTP errors should not trigger a pull attempt."""
-        hash_file, wordlist = self._setup_env(tmp_path)
-
-        mock_urlopen.side_effect = urllib.error.HTTPError(
-            url=f"{self.OLLAMA_URL}/api/generate",
-            code=500,
-            msg="Internal Server Error",
-            hdrs=None,
-            fp=None,
-        )
-
-        with mock.patch.object(hc_main, "ollamaUrl", self.OLLAMA_URL), \
-             mock.patch.object(hc_main, "ollamaModel", self.MODEL), \
-             mock.patch.object(hc_main, "hate_path", str(tmp_path)):
-
-            hc_main.hcatOllama("0", hash_file, "wordlist", wordlist)
-
-        captured = capsys.readouterr()
-        # HTTPError is a subclass of URLError, so it hits the URLError handler
-        assert "Could not connect to Ollama" in captured.out or "Error calling Ollama API" in captured.out
-        assert "500" in captured.out
-
-
-# ---------------------------------------------------------------------------
-# Step A: Mode routing, prompt construction, early-return error paths
-# ---------------------------------------------------------------------------
-
-class TestHcatOllamaModeRouting:
-    """Test mode selection, prompt building, and early-return errors."""
-
-    def test_unknown_mode_prints_error(self, ollama_env, capsys):
-        """Bad mode string → error message, no API call."""
-        with ollama_globals(ollama_env.tmp_path), \
-             mock.patch("hate_crack.main.urllib.request.urlopen") as mock_url:
-            hc_main.hcatOllama("0", ollama_env.hash_file, "bogus", "")
-
-        captured = capsys.readouterr()
-        assert "Unknown LLM generation mode" in captured.out
-        mock_url.assert_not_called()
-
-    def test_missing_wordlist_prints_error(self, ollama_env, capsys):
-        """Non-existent wordlist path → error, no API call."""
-        with ollama_globals(ollama_env.tmp_path), \
-             mock.patch("hate_crack.main.urllib.request.urlopen") as mock_url:
-            hc_main.hcatOllama(
-                "0", ollama_env.hash_file, "wordlist",
-                "/no/such/wordlist.txt",
-            )
-
-        captured = capsys.readouterr()
-        assert "Wordlist not found" in captured.out
-        mock_url.assert_not_called()
-
-    def test_wordlist_mode_reads_all_lines(self, ollama_env, capsys):
-        """All non-blank lines from the wordlist should appear in the prompt."""
-        # Write 600 lines to the wordlist
-        big_wordlist = ollama_env.tmp_path / "big.txt"
-        big_wordlist.write_text("\n".join(f"pass{i}" for i in range(600)) + "\n")
-
-        captured_payload = {}
-
-        def fake_urlopen(req, **kwargs):
-            captured_payload["data"] = json.loads(req.data.decode())
-            return _generate_response(["Password1"])
-
-        with ollama_globals(ollama_env.tmp_path), \
-             mock.patch("hate_crack.main.urllib.request.urlopen", side_effect=fake_urlopen), \
-             mock.patch("subprocess.Popen") as mock_popen:
-            proc = _make_proc()
-            mock_popen.return_value = proc
-            hc_main.hcatOllama(
-                "0", ollama_env.hash_file, "wordlist", str(big_wordlist),
-            )
-
-        prompt_text = captured_payload["data"]["prompt"]
-        # All lines should be included in the prompt
-        assert "pass0" in prompt_text
-        assert "pass499" in prompt_text
-        assert "pass599" in prompt_text
-
-    def test_target_mode_includes_context_in_prompt(self, ollama_env, capsys):
-        """Company, industry, and location should appear in the prompt."""
-        captured_payload = {}
-
-        def fake_urlopen(req, **kwargs):
-            captured_payload["data"] = json.loads(req.data.decode())
-            return _generate_response(["AcmeCorp2024"])
-
-        target_info = {
-            "company": "AcmeCorp",
-            "industry": "Finance",
-            "location": "New York",
-        }
-        with ollama_globals(ollama_env.tmp_path), \
-             mock.patch("hate_crack.main.urllib.request.urlopen", side_effect=fake_urlopen), \
-             mock.patch("subprocess.Popen") as mock_popen:
-            proc = _make_proc()
-            mock_popen.return_value = proc
-            hc_main.hcatOllama(
-                "0", ollama_env.hash_file, "target", target_info,
-            )
-
-        prompt_text = captured_payload["data"]["prompt"]
-        assert "AcmeCorp" in prompt_text
-        assert "Finance" in prompt_text
-        assert "New York" in prompt_text
-
-
-# ---------------------------------------------------------------------------
-# Step B: Candidate filtering / post-processing
-# ---------------------------------------------------------------------------
-
-class TestHcatOllamaCandidateFiltering:
-    """Test regex stripping, blank-line removal, 128-char limit, empty-result handling."""
-
-    def _run_with_response(self, ollama_env, response_text):
-        """Run hcatOllama with a canned Ollama response_text, return candidates file content."""
-        body = json.dumps({"response": response_text}).encode()
-        resp = mock.MagicMock()
-        resp.__enter__ = mock.Mock(return_value=io.BytesIO(body))
-        resp.__exit__ = mock.Mock(return_value=False)
-
-        candidates_path = f"{ollama_env.hash_file}.ollama_candidates"
-        captured = {}
-
-        def capture_popen(cmd, **kwargs):
-            # Read the candidates file before hashcat "runs" (cleanup happens after)
-            if os.path.isfile(candidates_path):
-                with open(candidates_path) as f:
-                    captured["content"] = f.read()
-            return _make_proc()
-
-        with ollama_globals(ollama_env.tmp_path), \
-             mock.patch("hate_crack.main.urllib.request.urlopen", return_value=resp), \
-             mock.patch("subprocess.Popen", side_effect=capture_popen):
-            hc_main.hcatOllama(
-                "0", ollama_env.hash_file, "wordlist", ollama_env.wordlist,
-            )
-
-        return captured.get("content")
-
-    def test_strips_numeric_dot_prefix(self, ollama_env):
-        content = self._run_with_response(ollama_env, "1. Password1\n2. Summer2024")
-        assert content is not None
-        lines = [l for l in content.strip().split("\n") if l]
-        assert "Password1" in lines
-        assert "Summer2024" in lines
-        # No line should start with a digit+dot
-        for line in lines:
-            assert not line.startswith("1.")
-            assert not line.startswith("2.")
-
-    def test_strips_dash_and_asterisk_prefix(self, ollama_env):
-        content = self._run_with_response(ollama_env, "- foo\n* bar")
-        assert content is not None
-        lines = [l for l in content.strip().split("\n") if l]
-        assert "foo" in lines
-        assert "bar" in lines
-
-    def test_skips_blank_lines(self, ollama_env):
-        content = self._run_with_response(ollama_env, "alpha\n\n\nbeta\n\n")
-        assert content is not None
-        lines = [l for l in content.strip().split("\n") if l]
-        assert lines == ["alpha", "beta"]
-
-    def test_rejects_over_128_chars(self, ollama_env):
-        long_pw = "A" * 129
-        content = self._run_with_response(ollama_env, f"short\n{long_pw}\nkeep")
-        assert content is not None
-        lines = [l for l in content.strip().split("\n") if l]
-        assert "short" in lines
-        assert "keep" in lines
-        assert long_pw not in lines
-
-    def test_accepts_exactly_128_chars(self, ollama_env):
-        exact_pw = "B" * 128
-        content = self._run_with_response(ollama_env, f"{exact_pw}\nother")
-        assert content is not None
-        lines = [l for l in content.strip().split("\n") if l]
-        assert exact_pw in lines
-
-    def test_empty_response_prints_error(self, ollama_env, capsys):
-        self._run_with_response(ollama_env, "")
-        captured = capsys.readouterr()
-        assert "no usable" in captured.out.lower()
-
-    def test_missing_response_key(self, ollama_env, capsys):
-        """If the JSON has no 'response' key, treat as empty."""
-        body = json.dumps({"other": "stuff"}).encode()
-        resp = mock.MagicMock()
-        resp.__enter__ = mock.Mock(return_value=io.BytesIO(body))
-        resp.__exit__ = mock.Mock(return_value=False)
-
-        with ollama_globals(ollama_env.tmp_path), \
-             mock.patch("hate_crack.main.urllib.request.urlopen", return_value=resp):
-            hc_main.hcatOllama(
-                "0", ollama_env.hash_file, "wordlist", ollama_env.wordlist,
-            )
-
-        captured = capsys.readouterr()
-        assert "no usable" in captured.out.lower()
-
-
-# ---------------------------------------------------------------------------
-# Step B: API error paths (non-404)
-# ---------------------------------------------------------------------------
-
-class TestHcatOllamaApiErrors:
-    """Test connection errors and generic exceptions during the generate call."""
-
-    def test_url_error_prints_connection_error(self, ollama_env, capsys):
-        with ollama_globals(ollama_env.tmp_path), \
-             mock.patch(
-                 "hate_crack.main.urllib.request.urlopen",
-                 side_effect=urllib.error.URLError("Connection refused"),
-             ):
-            hc_main.hcatOllama(
-                "0", ollama_env.hash_file, "wordlist", ollama_env.wordlist,
-            )
-
-        captured = capsys.readouterr()
-        assert "Could not connect" in captured.out
-        assert "Ensure Ollama is running" in captured.out
-
-    def test_generic_exception_prints_error(self, ollama_env, capsys):
-        with ollama_globals(ollama_env.tmp_path), \
-             mock.patch(
-                 "hate_crack.main.urllib.request.urlopen",
-                 side_effect=RuntimeError("boom"),
-             ):
-            hc_main.hcatOllama(
-                "0", ollama_env.hash_file, "wordlist", ollama_env.wordlist,
-            )
-
-        captured = capsys.readouterr()
-        assert "Error calling Ollama API" in captured.out
-
-    def test_generate_request_payload(self, ollama_env):
-        """Verify the /api/generate request has correct URL, model, stream=false."""
-        captured_req = {}
-
-        def fake_urlopen(req, **kwargs):
-            captured_req["url"] = req.full_url
-            captured_req["body"] = json.loads(req.data.decode())
-            return _generate_response(["Password1"])
-
-        with ollama_globals(ollama_env.tmp_path), \
-             mock.patch("hate_crack.main.urllib.request.urlopen", side_effect=fake_urlopen), \
-             mock.patch("subprocess.Popen") as mock_popen:
-            proc = _make_proc()
-            mock_popen.return_value = proc
-            hc_main.hcatOllama(
-                "0", ollama_env.hash_file, "wordlist", ollama_env.wordlist,
-            )
-
-        assert captured_req["url"] == f"{OLLAMA_URL}/api/generate"
-        assert captured_req["body"]["model"] == MODEL
-        assert captured_req["body"]["stream"] is False
-        assert len(captured_req["body"]["prompt"]) > 0
-
-
-# ---------------------------------------------------------------------------
-# Step C: Hashcat command construction
-# ---------------------------------------------------------------------------
-
-class TestHcatOllamaHashcatCommand:
-    """Test hashcat command flags and process handling."""
-
-    def _run_full(self, ollama_env, tuning="", potfile=""):
-        """Run hcatOllama through all steps, returning the list of Popen calls."""
-        popen_calls = []
-
-        def track_popen(cmd, **kwargs):
-            popen_calls.append((list(cmd), dict(kwargs)))
-            return _make_proc()
-
-        with ollama_globals(ollama_env.tmp_path, tuning=tuning, potfile=potfile), \
-             _urlopen_with_response(["Password1"]), \
-             mock.patch("subprocess.Popen", side_effect=track_popen), \
-             mock.patch("hate_crack.main.generate_session_id", return_value="test_session"):
-            hc_main.hcatOllama(
-                "1000", ollama_env.hash_file, "wordlist", ollama_env.wordlist,
-            )
-
-        return popen_calls
-
-    def test_hashcat_command_structure(self, ollama_env):
-        """First Popen call should be a straight wordlist attack with candidates file,
-        followed by rule-based attacks."""
-        calls = self._run_full(ollama_env)
-        assert len(calls) >= 1, "Expected at least one Popen call (hashcat)"
-
-        # First call: base wordlist attack (no rules)
-        hashcat_cmd = calls[0][0]
-        assert hashcat_cmd[0] == "/usr/bin/hashcat"
-        assert "-m" in hashcat_cmd
-        m_idx = hashcat_cmd.index("-m")
-        assert hashcat_cmd[m_idx + 1] == "1000"
-        # Should use the candidates wordlist file
-        candidates_path = f"{ollama_env.hash_file}.ollama_candidates"
-        assert candidates_path in hashcat_cmd
-        # First call should NOT have mask attack flags
-        assert "-a" not in hashcat_cmd
-        assert "--markov-hcstat2" not in " ".join(hashcat_cmd)
-        assert "--increment" not in hashcat_cmd
-        # First call should NOT have -r (rule) flag
-        assert "-r" not in hashcat_cmd
-
-        # Subsequent calls: rule-based attacks
-        if len(calls) > 1:
-            for rule_cmd, _ in calls[1:]:
-                assert "-r" in rule_cmd
-                assert candidates_path in rule_cmd
-
-    def test_hashcat_includes_tuning(self, ollama_env):
-        """-w 3 tuning flag should be appended to hashcat cmd."""
-        calls = self._run_full(ollama_env, tuning="-w 3")
-        hashcat_cmd = calls[0][0]
-        assert "-w" in hashcat_cmd
-        w_idx = hashcat_cmd.index("-w")
-        assert hashcat_cmd[w_idx + 1] == "3"
-
-    def test_hashcat_includes_potfile(self, ollama_env):
-        """--potfile-path should be present when hcatPotfilePath is set."""
-        calls = self._run_full(ollama_env, potfile="/tmp/test.potfile")
-        hashcat_cmd = calls[0][0]
-        potfile_flags = [f for f in hashcat_cmd if "--potfile-path=" in f]
-        assert len(potfile_flags) == 1
-        assert potfile_flags[0] == "--potfile-path=/tmp/test.potfile"
-
-    def test_hashcat_keyboard_interrupt(self, ollama_env, capsys):
-        """KeyboardInterrupt during hashcat should kill the process."""
-        proc = mock.MagicMock()
-        proc.pid = 99999
-        proc.wait.side_effect = KeyboardInterrupt()
-
-        with ollama_globals(ollama_env.tmp_path), \
-             _urlopen_with_response(["Password1"]), \
-             mock.patch("subprocess.Popen", return_value=proc), \
-             mock.patch("hate_crack.main.generate_session_id", return_value="test_session"):
-            hc_main.hcatOllama(
-                "0", ollama_env.hash_file, "wordlist", ollama_env.wordlist,
-            )
-
-        captured = capsys.readouterr()
-        assert "Killing PID" in captured.out
-        proc.kill.assert_called_once()
-
-
-# ---------------------------------------------------------------------------
-# Step F: Cleanup
-# ---------------------------------------------------------------------------
-
-class TestHcatOllamaWordlistPersistence:
-    """Test that the generated wordlist file persists after the run."""
-
-    def test_candidates_file_persists(self, ollama_env):
-        """ollama_candidates wordlist should remain after the run."""
-        with ollama_globals(ollama_env.tmp_path), \
-             _urlopen_with_response(["Password1"]), \
-             mock.patch("subprocess.Popen", return_value=_make_proc()), \
-             mock.patch("hate_crack.main.generate_session_id", return_value="test_session"):
-            hc_main.hcatOllama(
-                "0", ollama_env.hash_file, "wordlist", ollama_env.wordlist,
-            )
-
-        candidates_path = f"{ollama_env.hash_file}.ollama_candidates"
-        assert os.path.isfile(candidates_path), "candidates wordlist should persist"
-
-
-# ---------------------------------------------------------------------------
-# attacks.py: ollama_attack() UI handler
-# ---------------------------------------------------------------------------
-
-class TestOllamaAttackHandler:
-    """Test the ollama_attack(ctx) menu handler in attacks.py."""
-
-    def _make_ctx(self, **overrides):
-        """Build a mock ctx with the attributes ollama_attack() reads."""
-        ctx = mock.MagicMock()
-        ctx.hcatHashType = "0"
-        ctx.hcatHashFile = "/tmp/hashes.txt"
-        ctx.hcatWordlists = "/tmp/wordlists"
-        for k, v in overrides.items():
-            setattr(ctx, k, v)
-        return ctx
-
-    def test_target_mode(self):
-        """ollama_attack prompts for company/industry/location and calls hcatOllama."""
-        from hate_crack.attacks import ollama_attack
-
-        ctx = self._make_ctx()
-        with mock.patch("builtins.input", side_effect=["AcmeCorp", "Finance", "NYC"]):
-            ollama_attack(ctx)
-
-        ctx.hcatOllama.assert_called_once()
-        call_args = ctx.hcatOllama.call_args
-        assert call_args[0][2] == "target"
-        target_info = call_args[0][3]
-        assert target_info["company"] == "AcmeCorp"
-        assert target_info["industry"] == "Finance"
-        assert target_info["location"] == "NYC"
-