feat(llm): structured candidate generation module via Atomic Agents

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

hate_crack/llm.py

@@ -0,0 +1,134 @@
+"""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``.
+"""
+
+from typing import Any
+
+import instructor
+from openai import OpenAI
+from pydantic import Field
+
+from atomic_agents import AgentConfig, AtomicAgent, BaseIOSchema
+from atomic_agents.context import SystemPromptGenerator
+
+MAX_CANDIDATE_LEN = 128
+
+
+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,
+) -> list[str]:
+    """Generate password candidates via an Ollama-backed AtomicAgent.
+
+    Returns a deduped, length-capped list of candidate strings (may be empty).
+    Raises ValueError for an unknown mode. 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"),
+        mode=instructor.Mode.JSON,
+    )
+    prompt_generator = _TARGET_PROMPT if mode == "target" else _WORDLIST_PROMPT
+
+    agent = AtomicAgent[GenerationInput, PasswordCandidatesOutput](
+        config=AgentConfig.model_construct(
+            client=client,
+            model=model,
+            system_prompt_generator=prompt_generator,
+            model_api_parameters={"extra_body": {"options": {"num_ctx": num_ctx}}},
+        )
+    )
+
+    result: Any = agent.run(GenerationInput(request=request))
+
+    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

tests/test_llm.py

@@ -0,0 +1,100 @@
+"""Unit tests for hate_crack.llm.generate_candidates."""
+
+import os
+from unittest import mock
+
+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"),
+        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_unknown_mode_raises():
+    with pytest.raises(ValueError):
+        llm.generate_candidates(
+            "http://localhost:11434", "qwen2.5:32b", 2048, "bogus", {},
+        )