diff --git a/README.md b/README.md index 6b10e8c..32e2738 100644 --- a/README.md +++ b/README.md @@ -473,11 +473,33 @@ The LLM Attack (option 12) uses Ollama to generate password candidates. Configur - **`ollamaNumCtx`** — Context window size for the model (default: `2048`). - **`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. - **`ollamaMaxSampleLines`** — Maximum number of lines drawn from the source file and included in the LLM prompt when using **Wordlist** or **Cracked passwords** mode (default: `500`). Lines are sampled evenly across the whole file so the prompt reflects the wordlist's full character range rather than just the head. Set to a larger value if the model has a big context window (`ollamaNumCtx`) and you want richer coverage; set it lower to reduce prompt size and generation latency. Values ≤ 0 are treated as 500. +- **`ollamaAutoResearch`** — When `true` (default), **Target info** mode asks the local model to suggest the industry and location as soon as you have typed the company name, and offers them as editable prompt defaults. Set to `false` to always get blank prompts (useful with a slow model, since research costs one extra round-trip before the attack starts). - 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 three generation modes: 1. **Target info** — company / industry / location; the model derives candidates from those details. + + After you type the company name, hate_crack asks the same local model what it already knows about that organization and pre-fills the **Industry** and **Location** prompts with the answers, shown in parentheses: + + ``` + Company name: Acme Rail Services + + [!] The values in parentheses below are the local model's GUESSES, not verified OSINT. + Press Enter to accept, or type your own value to override. + Industry (freight rail maintenance): + Location (Omaha, Nebraska): + ``` + + Press Enter to accept a suggestion or type over it. These values are the model's recollection, **not OSINT** — treat them as a starting point, not intelligence about the client. The lookup uses only the local Ollama server, so the client name never leaves the host; there are no web or third-party API calls. If the model does not recognize the organization (the common case for small clients), it returns nothing and you get plain blank prompts: + + ``` + Company name: Acme Rail Services + Industry: + Location: + ``` + + A research failure — timeout, Ollama not running, empty answer — never blocks the attack; it just falls back to blank prompts. Set `ollamaAutoResearch` to `false` to skip research entirely. 2. **Wordlist** — derive basewords from a sample wordlist. 3. **Cracked passwords** — feed the plaintexts already recovered this session (`.out`) back to the model so it can infer the target organization's own password conventions (basewords, seasons, years, suffixes, leetspeak) and generate *new* candidates in the same style. This option is only listed once at least one hash has been cracked; the sample is capped by `ollamaMaxSampleLines` exactly like Wordlist mode. diff --git a/config.json.example b/config.json.example index 86bd196..ec89e9a 100644 --- a/config.json.example +++ b/config.json.example @@ -29,6 +29,7 @@ "ollamaNumCtx": 2048, "ollamaTimeout": 300, "ollamaMaxSampleLines": 500, + "ollamaAutoResearch": true, "omenTrainingList": "rockyou.txt", "omenMaxCandidates": 50000000, "pcfgRuleset": "DEFAULT", diff --git a/hate_crack/attacks.py b/hate_crack/attacks.py index 4288d94..c84ccf2 100644 --- a/hate_crack/attacks.py +++ b/hate_crack/attacks.py @@ -7,6 +7,7 @@ from typing import Any from hate_crack import notify as _notify from hate_crack.api import download_hashmob_rules from hate_crack.formatting import print_multicolumn_list +from hate_crack.llm import clean_research_field from hate_crack.menu import interactive_menu @@ -510,6 +511,47 @@ def bandrel_method(ctx: Any) -> None: ctx.hcatBandrel(ctx.hcatHashType, ctx.hcatHashFile) +def _research_target_suggestions(ctx: Any, company: str) -> dict[str, str]: + """Ask the local model for industry/location suggestions for *company*. + + Returns a dict of cleaned suggestion strings (values may be ''). Research is + a convenience only, so any failure is swallowed here as well as in + ``hcatOllamaResearchTarget``: the operator still gets blank prompts and the + attack proceeds. + """ + if not company: + return {} + + try: + raw = ctx.hcatOllamaResearchTarget(company) + except Exception as e: + print(f"Note: target research unavailable ({e}) — enter the details manually.") + return {} + + suggestions = {} + if isinstance(raw, dict): + for key in ("industry", "location"): + value = clean_research_field(raw.get(key, "")) + if value: + suggestions[key] = value + + if suggestions: + print( + "\n[!] The values in parentheses below are the local model's GUESSES, " + "not verified OSINT." + ) + print(" Press Enter to accept, or type your own value to override.") + return suggestions + + +def _prompt_with_default(label: str, default: Any) -> str: + """Prompt for *label*, showing *default* in parentheses when there is one.""" + suggestion = clean_research_field(default) + if suggestion: + return input(f"{label} ({suggestion}): ").strip() or suggestion + return input(f"{label}: ").strip() + + def ollama_attack(ctx: Any) -> None: _notify.prompt_notify_for_attack("LLM") # Cracked-password mode is only offered when this session actually has @@ -531,8 +573,9 @@ def ollama_attack(ctx: Any) -> None: return if choice == "1": company = input("Company name: ").strip() - industry = input("Industry: ").strip() - location = input("Location: ").strip() + suggestions = _research_target_suggestions(ctx, company) + industry = _prompt_with_default("Industry", suggestions.get("industry")) + location = _prompt_with_default("Location", suggestions.get("location")) ctx.hcatOllama( ctx.hcatHashType, ctx.hcatHashFile, diff --git a/hate_crack/config.json.example b/hate_crack/config.json.example index 0bd02d7..fe8ffb8 100644 --- a/hate_crack/config.json.example +++ b/hate_crack/config.json.example @@ -26,6 +26,7 @@ "ollamaNumCtx": 2048, "ollamaTimeout": 300, "ollamaMaxSampleLines": 500, + "ollamaAutoResearch": true, "passgptModel": "javirandor/passgpt-10characters", "passgptMaxCandidates": 1000000, "passgptBatchSize": 1024, diff --git a/hate_crack/llm.py b/hate_crack/llm.py index f4cf5f5..9600489 100644 --- a/hate_crack/llm.py +++ b/hate_crack/llm.py @@ -1,7 +1,7 @@ """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``. +to this module only through ``generate_candidates`` and ``research_target``. """ import instructor @@ -14,6 +14,10 @@ from atomic_agents.context import SystemPromptGenerator MAX_CANDIDATE_LEN = 128 DEFAULT_TIMEOUT_SECONDS = 300.0 +# Researched target fields are pasted into an interactive prompt as an editable +# default, so they must stay short enough to fit on one terminal line. +MAX_RESEARCH_FIELD_LEN = 80 + class LLMTimeoutError(Exception): """The LLM server accepted the request but did not respond in time. @@ -43,6 +47,60 @@ class PasswordCandidatesOutput(BaseIOSchema): ) +class TargetResearchInput(BaseIOSchema): + """The company name to recall industry and location details for.""" + + company: str = Field( + ..., description="The name of the target organization, exactly as typed." + ) + + +class TargetResearchOutput(BaseIOSchema): + """Recalled industry and location for a named organization, or empty strings.""" + + industry: str = Field( + ..., + description=( + "The organization's industry or sector as a short phrase, e.g. " + "'regional healthcare provider' or 'commercial construction'. Return " + "an empty string if you do not actually recognize this organization." + ), + ) + location: str = Field( + ..., + description=( + "The organization's primary location as 'City, State/Country'. Return " + "an empty string if you do not actually recognize this organization." + ), + ) + + +_RESEARCH_PROMPT = SystemPromptGenerator( + background=[ + "You are assisting a security professional on an authorized penetration " + "test who is about to generate password candidates for a named client " + "organization.", + "You have no internet access. You may only answer from what you already " + "know about the organization.", + "Most client organizations are small and will be completely unknown to " + "you. That is the expected case, not a failure.", + ], + steps=[ + "Decide whether you genuinely recognize this specific organization by name.", + "If you do, recall its industry or sector and its primary location.", + "If you do not recognize it, or you are not reasonably confident, do not " + "guess and do not infer anything from the words in the name.", + ], + output_instructions=[ + "Return an empty string for any field you are not reasonably confident " + "about. An empty field is correct and useful; a fabricated one is harmful " + "because the operator may mistake it for real intelligence.", + "Keep each field under 80 characters.", + "Return only the industry and location fields — no explanations, " + "hedging, caveats, or commentary.", + ], +) + _TARGET_PROMPT = SystemPromptGenerator( background=[ "You are a security professional generating password candidates during an " @@ -135,6 +193,70 @@ def _build_request(mode: str, context_data: dict) -> str: raise ValueError(f"Unknown LLM generation mode: {mode}") +def _build_client(url: str, timeout: float) -> instructor.Instructor: + """Build the instructor-wrapped OpenAI client pointed at an Ollama server.""" + return instructor.from_openai( + OpenAI(base_url=f"{url}/v1", api_key="ollama", timeout=timeout), + mode=instructor.Mode.JSON, + ) + + +def clean_research_field(value: object) -> str: + """Strip and length-cap one researched field; return '' for no suggestion. + + Anything that is not a non-empty string after stripping — including a model + that echoed whitespace or an over-long ramble — becomes '' so the caller + falls back to a plain blank prompt instead of pasting model noise into it. + """ + if not isinstance(value, str): + return "" + cleaned = " ".join(value.split()) + if len(cleaned) > MAX_RESEARCH_FIELD_LEN: + cleaned = cleaned[:MAX_RESEARCH_FIELD_LEN].rstrip() + return cleaned + + +def research_target( + url: str, + model: str, + num_ctx: int, + company: str, + timeout: float = DEFAULT_TIMEOUT_SECONDS, +) -> TargetResearchOutput: + """Ask the local model what it already knows about *company*. + + Returns a ``TargetResearchOutput`` whose ``industry`` and ``location`` are + stripped and capped at ``MAX_RESEARCH_FIELD_LEN``; either may be '' when the + model is not confident, which callers must treat as "no suggestion". + + Uses only the configured local Ollama server — no web lookups, so the client + name never leaves the host. Raises LLMTimeoutError if the request exceeds + ``timeout``; other client/connection errors propagate to the caller. + """ + client = _build_client(url, timeout) + + agent = AtomicAgent[TargetResearchInput, TargetResearchOutput]( + config=AgentConfig( + client=client, + model=model, + system_prompt_generator=_RESEARCH_PROMPT, + model_api_parameters={"extra_body": {"options": {"num_ctx": num_ctx}}}, + ) + ) + + try: + result = agent.run(TargetResearchInput(company=company)) + except APITimeoutError as e: + raise LLMTimeoutError( + f"no response from {url} within {timeout:g} seconds" + ) from e + + return TargetResearchOutput( + industry=clean_research_field(getattr(result, "industry", "")), + location=clean_research_field(getattr(result, "location", "")), + ) + + def generate_candidates( url: str, model: str, @@ -156,10 +278,7 @@ def generate_candidates( """ 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, - ) + client = _build_client(url, timeout) # _build_request has already rejected unknown modes, so this lookup is safe. prompt_generator = _PROMPTS[mode] diff --git a/hate_crack/main.py b/hate_crack/main.py index b0818eb..f6dafd0 100755 --- a/hate_crack/main.py +++ b/hate_crack/main.py @@ -455,6 +455,7 @@ ollamaModel = config_parser.get("ollamaModel", "qwen2.5:32b") ollamaNumCtx = int(config_parser.get("ollamaNumCtx", 2048)) ollamaTimeout = float(config_parser.get("ollamaTimeout", 300)) ollamaMaxSampleLines = int(config_parser.get("ollamaMaxSampleLines", 500)) +ollamaAutoResearch = bool(config_parser.get("ollamaAutoResearch", True)) omenTrainingList = config_parser.get("omenTrainingList", "rockyou.txt") omenMaxCandidates = int(config_parser.get("omenMaxCandidates", 1000000)) @@ -2123,6 +2124,43 @@ def _sample_plaintext_file(path, cap, source_label="wordlist"): return sampled +def hcatOllamaResearchTarget(company): + """Ask the local Ollama model what it knows about *company*. + + Returns a dict with "industry" and "location" keys; either value may be an + empty string when the model is not confident or the request failed. Never + raises: research is a convenience, so any failure degrades to empty + suggestions (blank prompts) rather than blocking the attack. + + Uses only the configured local Ollama server — the company name is never + sent to a third-party service. + """ + blank = {"industry": "", "location": ""} + if not ollamaAutoResearch or not company: + return blank + + try: + with spinner(f"Researching {company} via Ollama ({ollamaModel})..."): + result = llm.research_target( + ollamaUrl, + ollamaModel, + ollamaNumCtx, + company, + timeout=ollamaTimeout, + ) + except llm.LLMTimeoutError: + print( + f"Note: target research timed out after {ollamaTimeout:g} seconds — " + "enter the details manually." + ) + return blank + except Exception as e: + print(f"Note: target research unavailable ({e}) — enter the details manually.") + return blank + + return {"industry": result.industry, "location": result.location} + + # LLM Ollama Attack def hcatOllama(hcatHashType, hcatHashFile, mode, context_data): candidates_path = f"{hcatHashFile}.ollama_candidates" diff --git a/tests/test_llm.py b/tests/test_llm.py index 70c58ed..1e8ef94 100644 --- a/tests/test_llm.py +++ b/tests/test_llm.py @@ -210,3 +210,113 @@ def test_api_timeout_reraised_as_domain_error(): "target", {"company": "X", "industry": "Y", "location": "Z"}, timeout=1.0, ) + + +# --------------------------------------------------------------------------- +# research_target +# --------------------------------------------------------------------------- + + +def _patch_research_agent(industry, location): + """Patch client builders + AtomicAgent for a research call. No network.""" + result = mock.MagicMock() + result.industry = industry + result.location = location + + agent_instance = mock.MagicMock() + agent_instance.run.return_value = result + + agent_cls = mock.MagicMock() + 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_research_target_returns_fields_and_passes_company(): + p_instr, p_openai, p_agent, agent_cls, agent_instance = _patch_research_agent( + "freight rail maintenance", "Omaha, Nebraska" + ) + with p_instr, p_openai, p_agent: + out = llm.research_target( + "http://localhost:11434", "qwen2.5:32b", 2048, "Acme Rail Services" + ) + assert out.industry == "freight rail maintenance" + assert out.location == "Omaha, Nebraska" + assert agent_instance.run.call_args[0][0].company == "Acme Rail Services" + + +def test_research_target_uses_research_prompt_and_num_ctx(): + p_instr, p_openai, p_agent, agent_cls, agent_instance = _patch_research_agent( + "x", "y" + ) + with p_instr, p_openai, p_agent: + llm.research_target("http://localhost:11434", "qwen2.5:32b", 4096, "Acme") + config = agent_cls.__getitem__.return_value.call_args.kwargs["config"] + assert config.system_prompt_generator is llm._RESEARCH_PROMPT + assert config.model_api_parameters["extra_body"]["options"]["num_ctx"] == 4096 + + +def test_research_prompt_tells_model_to_return_empty_when_unsure(): + rendered = llm._RESEARCH_PROMPT.generate_prompt().lower() + assert "empty string" in rendered + assert "no internet access" in rendered + + +def test_research_target_strips_and_blanks_whitespace_only(): + p_instr, p_openai, p_agent, agent_cls, agent_instance = _patch_research_agent( + " healthcare ", " " + ) + with p_instr, p_openai, p_agent: + out = llm.research_target("http://localhost:11434", "m", 2048, "Acme") + assert out.industry == "healthcare" + assert out.location == "" + + +def test_research_target_caps_overlong_values(): + p_instr, p_openai, p_agent, agent_cls, agent_instance = _patch_research_agent( + "A" * 500, "B" * (llm.MAX_RESEARCH_FIELD_LEN + 1) + ) + with p_instr, p_openai, p_agent: + out = llm.research_target("http://localhost:11434", "m", 2048, "Acme") + assert len(out.industry) == llm.MAX_RESEARCH_FIELD_LEN + assert len(out.location) == llm.MAX_RESEARCH_FIELD_LEN + + +def test_research_target_tolerates_non_string_fields(): + p_instr, p_openai, p_agent, agent_cls, agent_instance = _patch_research_agent( + None, 42 + ) + with p_instr, p_openai, p_agent: + out = llm.research_target("http://localhost:11434", "m", 2048, "Acme") + assert out.industry == "" + assert out.location == "" + + +def test_research_target_timeout_forwarded_and_translated(): + p_instr, p_openai, p_agent, agent_cls, agent_instance = _patch_research_agent( + "x", "y" + ) + agent_instance.run.side_effect = openai.APITimeoutError( + request=httpx.Request("POST", "http://localhost:11434/v1/chat/completions") + ) + with p_instr, p_openai as openai_cls, p_agent: + with pytest.raises(llm.LLMTimeoutError): + llm.research_target( + "http://localhost:11434", "m", 2048, "Acme", timeout=7.5 + ) + assert openai_cls.call_args.kwargs["timeout"] == 7.5 + + +def test_clean_research_field_collapses_internal_whitespace(): + assert llm.clean_research_field("commercial \n construction") == ( + "commercial construction" + ) diff --git a/tests/test_ollama_target_research.py b/tests/test_ollama_target_research.py new file mode 100644 index 0000000..222240a --- /dev/null +++ b/tests/test_ollama_target_research.py @@ -0,0 +1,301 @@ +"""Tests for LLM target-research pre-fill (menu 12, Target info mode). + +Covers both layers: hcatOllamaResearchTarget in main (spinner + failure +degradation) and ollama_attack's editable-default prompts in attacks. The LLM is +always mocked; no network. +""" + +import os +from unittest import mock +from unittest.mock import MagicMock, patch + +os.environ["HATE_CRACK_SKIP_INIT"] = "1" +from hate_crack import llm # noqa: E402 +from hate_crack import main as hc_main # noqa: E402 +from hate_crack.attacks import ollama_attack # noqa: E402 + + +def _make_ctx(hash_type: str = "1000", hash_file: str = "/tmp/hashes.txt") -> MagicMock: + ctx = MagicMock() + ctx.hcatHashType = hash_type + ctx.hcatHashFile = hash_file + return ctx + + +class _InputRecorder: + """input() stub that records prompts and replays canned answers.""" + + def __init__(self, answers): + self.answers = list(answers) + self.prompts: list[str] = [] + + def __call__(self, prompt=""): + self.prompts.append(prompt) + return self.answers.pop(0) + + +def _run_target_mode(ctx, answers): + recorder = _InputRecorder(answers) + with ( + patch("hate_crack.attacks.interactive_menu", return_value="1"), + patch("builtins.input", recorder), + ): + ollama_attack(ctx) + return recorder + + +# --------------------------------------------------------------------------- +# attacks layer: prompt defaults +# --------------------------------------------------------------------------- + + +class TestResearchPrefill: + def test_successful_research_prefills_defaults(self) -> None: + ctx = _make_ctx() + ctx.hcatOllamaResearchTarget.return_value = { + "industry": "freight rail maintenance", + "location": "Omaha, Nebraska", + } + + recorder = _run_target_mode(ctx, ["Acme Rail", "", ""]) + + ctx.hcatOllamaResearchTarget.assert_called_once_with("Acme Rail") + assert recorder.prompts == [ + "Company name: ", + "Industry (freight rail maintenance): ", + "Location (Omaha, Nebraska): ", + ] + + def test_enter_accepts_the_suggested_defaults(self) -> None: + ctx = _make_ctx() + ctx.hcatOllamaResearchTarget.return_value = { + "industry": "healthcare", + "location": "Austin, Texas", + } + + _run_target_mode(ctx, ["Acme", "", ""]) + + assert ctx.hcatOllama.call_args[0][3] == { + "company": "Acme", + "industry": "healthcare", + "location": "Austin, Texas", + } + + def test_typed_input_overrides_the_default(self) -> None: + ctx = _make_ctx() + ctx.hcatOllamaResearchTarget.return_value = { + "industry": "healthcare", + "location": "Austin, Texas", + } + + _run_target_mode(ctx, ["Acme", " banking ", "Berlin"]) + + assert ctx.hcatOllama.call_args[0][3] == { + "company": "Acme", + "industry": "banking", + "location": "Berlin", + } + + def test_partial_research_prefills_only_known_field(self) -> None: + ctx = _make_ctx() + ctx.hcatOllamaResearchTarget.return_value = { + "industry": "mining", + "location": "", + } + + recorder = _run_target_mode(ctx, ["Acme", "", "Perth"]) + + assert recorder.prompts[1] == "Industry (mining): " + assert recorder.prompts[2] == "Location: " + assert ctx.hcatOllama.call_args[0][3]["location"] == "Perth" + + def test_empty_research_falls_back_to_blank_prompts(self) -> None: + ctx = _make_ctx() + ctx.hcatOllamaResearchTarget.return_value = {"industry": "", "location": ""} + + recorder = _run_target_mode(ctx, ["Acme", "tech", "NYC"]) + + assert recorder.prompts == ["Company name: ", "Industry: ", "Location: "] + assert ctx.hcatOllama.call_args[0][3] == { + "company": "Acme", + "industry": "tech", + "location": "NYC", + } + + def test_whitespace_only_research_is_treated_as_no_suggestion(self) -> None: + ctx = _make_ctx() + ctx.hcatOllamaResearchTarget.return_value = { + "industry": " ", + "location": "\t\n", + } + + recorder = _run_target_mode(ctx, ["Acme", "", ""]) + + assert recorder.prompts == ["Company name: ", "Industry: ", "Location: "] + + def test_overlong_suggestion_is_capped_in_the_prompt(self) -> None: + ctx = _make_ctx() + ctx.hcatOllamaResearchTarget.return_value = { + "industry": "z" * 500, + "location": "", + } + + recorder = _run_target_mode(ctx, ["Acme", "", ""]) + + assert recorder.prompts[1] == f"Industry ({'z' * llm.MAX_RESEARCH_FIELD_LEN}): " + industry = ctx.hcatOllama.call_args[0][3]["industry"] + assert len(industry) == llm.MAX_RESEARCH_FIELD_LEN + + def test_research_exception_falls_back_and_does_not_abort(self) -> None: + """Even an unexpected raise from the research call must not block.""" + ctx = _make_ctx() + ctx.hcatOllamaResearchTarget.side_effect = RuntimeError("boom") + + recorder = _run_target_mode(ctx, ["Acme", "tech", "NYC"]) + + assert recorder.prompts == ["Company name: ", "Industry: ", "Location: "] + ctx.hcatOllama.assert_called_once() + + def test_non_dict_research_result_falls_back(self) -> None: + ctx = _make_ctx() + ctx.hcatOllamaResearchTarget.return_value = None + + recorder = _run_target_mode(ctx, ["Acme", "tech", "NYC"]) + + assert recorder.prompts == ["Company name: ", "Industry: ", "Location: "] + + def test_blank_company_skips_research_entirely(self) -> None: + ctx = _make_ctx() + + recorder = _run_target_mode(ctx, ["", "tech", "NYC"]) + + ctx.hcatOllamaResearchTarget.assert_not_called() + assert recorder.prompts == ["Company name: ", "Industry: ", "Location: "] + + def test_suggestions_are_labelled_as_guesses(self, capsys) -> None: + ctx = _make_ctx() + ctx.hcatOllamaResearchTarget.return_value = { + "industry": "healthcare", + "location": "Austin, Texas", + } + + _run_target_mode(ctx, ["Acme", "", ""]) + + out = capsys.readouterr().out + assert "GUESSES" in out + assert "not verified OSINT" in out + + def test_no_guess_banner_when_research_returns_nothing(self, capsys) -> None: + ctx = _make_ctx() + ctx.hcatOllamaResearchTarget.return_value = {"industry": "", "location": ""} + + _run_target_mode(ctx, ["Acme", "", ""]) + + assert "GUESSES" not in capsys.readouterr().out + + def test_hcatOllama_call_shape_unchanged(self) -> None: + ctx = _make_ctx(hash_type="1800", hash_file="/tmp/sha512.txt") + ctx.hcatOllamaResearchTarget.return_value = { + "industry": "healthcare", + "location": "Austin", + } + + _run_target_mode(ctx, ["Acme", "", ""]) + + args = ctx.hcatOllama.call_args[0] + assert args[0] == "1800" + assert args[1] == "/tmp/sha512.txt" + assert args[2] == "target" + assert set(args[3]) == {"company", "industry", "location"} + + +# --------------------------------------------------------------------------- +# main layer: hcatOllamaResearchTarget +# --------------------------------------------------------------------------- + + +class TestHcatOllamaResearchTarget: + def test_returns_researched_fields_and_shows_progress(self) -> None: + result = hc_main.llm.TargetResearchOutput(industry="mining", location="Perth") + with ( + mock.patch.object(hc_main, "ollamaAutoResearch", True), + mock.patch.object(hc_main, "ollamaModel", "qwen2.5:32b"), + mock.patch.object(hc_main, "ollamaTimeout", 30.0), + mock.patch.object(hc_main.llm, "research_target", return_value=result), + mock.patch.object(hc_main, "spinner") as spin, + ): + out = hc_main.hcatOllamaResearchTarget("Acme") + + assert out == {"industry": "mining", "location": "Perth"} + # The call must be wrapped in the progress spinner, not look like a hang. + assert spin.call_count == 1 + assert "Acme" in spin.call_args[0][0] + + def test_forwards_config_values_to_research_target(self) -> None: + result = hc_main.llm.TargetResearchOutput(industry="", location="") + with ( + mock.patch.object(hc_main, "ollamaAutoResearch", True), + mock.patch.object(hc_main, "ollamaUrl", "http://localhost:11434"), + mock.patch.object(hc_main, "ollamaModel", "m"), + mock.patch.object(hc_main, "ollamaNumCtx", 4096), + mock.patch.object(hc_main, "ollamaTimeout", 12.5), + mock.patch.object( + hc_main.llm, "research_target", return_value=result + ) as research, + ): + hc_main.hcatOllamaResearchTarget("Acme") + + assert research.call_args[0] == ("http://localhost:11434", "m", 4096, "Acme") + assert research.call_args.kwargs["timeout"] == 12.5 + + def test_timeout_degrades_to_blank_suggestions(self, capsys) -> None: + with ( + mock.patch.object(hc_main, "ollamaAutoResearch", True), + mock.patch.object(hc_main, "ollamaTimeout", 5.0), + mock.patch.object( + hc_main.llm, + "research_target", + side_effect=hc_main.llm.LLMTimeoutError("no response"), + ), + ): + out = hc_main.hcatOllamaResearchTarget("Acme") + + assert out == {"industry": "", "location": ""} + assert "timed out" in capsys.readouterr().out + + def test_connection_failure_degrades_to_blank_suggestions(self, capsys) -> None: + with ( + mock.patch.object(hc_main, "ollamaAutoResearch", True), + mock.patch.object( + hc_main.llm, + "research_target", + side_effect=ConnectionError("connection refused"), + ), + ): + out = hc_main.hcatOllamaResearchTarget("Acme") + + assert out == {"industry": "", "location": ""} + assert "unavailable" in capsys.readouterr().out + + def test_disabled_by_config_skips_the_model_call(self) -> None: + with ( + mock.patch.object(hc_main, "ollamaAutoResearch", False), + mock.patch.object(hc_main.llm, "research_target") as research, + ): + out = hc_main.hcatOllamaResearchTarget("Acme") + + research.assert_not_called() + assert out == {"industry": "", "location": ""} + + def test_blank_company_skips_the_model_call(self) -> None: + with ( + mock.patch.object(hc_main, "ollamaAutoResearch", True), + mock.patch.object(hc_main.llm, "research_target") as research, + ): + out = hc_main.hcatOllamaResearchTarget("") + + research.assert_not_called() + assert out == {"industry": "", "location": ""} + + def test_auto_research_default_is_enabled(self) -> None: + assert hc_main.ollamaAutoResearch is True