diff --git a/hate_crack/llm.py b/hate_crack/llm.py new file mode 100644 index 0000000..195e5d8 --- /dev/null +++ b/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 diff --git a/tests/test_llm.py b/tests/test_llm.py new file mode 100644 index 0000000..f96b39c --- /dev/null +++ b/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=) — 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", {}, + )