diff --git a/cyberai/agents/recon/agent.py b/cyberai/agents/recon/agent.py index 5017641..57ce001 100644 --- a/cyberai/agents/recon/agent.py +++ b/cyberai/agents/recon/agent.py @@ -8,8 +8,9 @@ from cyberai.core.scan_session import Severity from cyberai.core.types import OpenPort, ReconResult -from .dns_tool import detect_subdomains, run_dns, run_whois +from .dns_tool import run_dns, run_whois from .llm_detector import detect_llm_endpoints +from .subdomain_enum import enumerate_subdomains, fqdns from .web_surface import discover_surface from .behavioral import BehavioralFingerprint from .behavioral_probe import build_probe_context @@ -54,7 +55,7 @@ def _register_tools(self) -> None: Tool( name="subdomain_scan", description="Subdomain bruteforce", - func=detect_subdomains, + func=enumerate_subdomains, parameters={"target": "str"}, ) ) @@ -113,7 +114,7 @@ def run(self, target: str, context: Optional[Dict[str, Any]] = None) -> Dict[str # 4. Subdomains self._check_iteration_limit() - sub_result = detect_subdomains(target) + sub_result = enumerate_subdomains(target) self.kb.set("recon.subdomains", sub_result, agent=self.AGENT_NAME) results["recon.subdomains"] = sub_result self._log("subdomain_scan complete", sub_result) @@ -161,7 +162,7 @@ def run(self, target: str, context: Optional[Dict[str, Any]] = None) -> Dict[str ports=[OpenPort(**p) for p in ports if isinstance(p, dict)], whois=whois_result if isinstance(whois_result, dict) else {}, dns=dns_result if isinstance(dns_result, dict) else {}, - subdomains=(sub_result.get("subdomains", []) if isinstance(sub_result, dict) else []), + subdomains=fqdns(sub_result if isinstance(sub_result, dict) else None), ) self.kb.set("recon.result", recon_result.model_dump(), agent=self.AGENT_NAME) diff --git a/cyberai/agents/recon/subdomain_enum.py b/cyberai/agents/recon/subdomain_enum.py index 17ed27a..deafde2 100644 --- a/cyberai/agents/recon/subdomain_enum.py +++ b/cyberai/agents/recon/subdomain_enum.py @@ -64,6 +64,18 @@ ] +def fqdns(result: Dict[str, Any] | None = None) -> List[str]: + """Extract plain hostnames from an enumerate_subdomains* result. + + Both enumerators return {"found": [{"fqdn": ...}]}. Callers that need + bare hostnames (ReconResult.subdomains, the planner KB graph) must use + this instead of re-deriving the shape, or the sync and async pipeline + paths drift apart silently. + """ + found = (result or {}).get("found") or [] + return [r["fqdn"] for r in found if isinstance(r, dict) and r.get("fqdn")] + + def enumerate_subdomains( domain: str, wordlist: List[str] = None, diff --git a/cyberai/core/llm_client.py b/cyberai/core/llm_client.py index fac4322..3471757 100644 --- a/cyberai/core/llm_client.py +++ b/cyberai/core/llm_client.py @@ -10,6 +10,11 @@ from .base_agent import Tool +# Local models on a 3060 are slow on long exploit prompts; the async path +# used to allow only 60s and timed out where the sync path succeeded. +OLLAMA_TIMEOUT = 120 + + class LLMClient: """ Unified LLM interface — OpenAI / Anthropic / Ollama @@ -122,9 +127,14 @@ def _call_anthropic( ) return response.content[0].text - def _call_ollama( - self, messages: List[Dict], system: Optional[str], agent_name: str = "unknown" - ) -> str: + def _ollama_request(self, messages: List[Dict], system: Optional[str]) -> tuple: + """Build the (url, payload) pair for an ollama /api/chat call. + + Single source for both the sync and async entry points: the async one + used to build its own payload and had silently dropped the system + prompt and the raised num_ctx, so the same call behaved differently + depending on which path reached it. + """ url = f"{self.config.base_url or 'http://localhost:11434'}/api/chat" full_messages: List[Dict] = [] if system: @@ -138,7 +148,13 @@ def _call_ollama( # prompts (CVE JSON + attack paths + chain), which 4xx/5xx the call. "options": {"num_ctx": 8192}, } - response = httpx.post(url, json=payload, timeout=120) + return url, payload + + def _call_ollama( + self, messages: List[Dict], system: Optional[str], agent_name: str = "unknown" + ) -> str: + url, payload = self._ollama_request(messages, system) + response = httpx.post(url, json=payload, timeout=OLLAMA_TIMEOUT) if response.status_code != 200: raise RuntimeError(f"ollama HTTP {response.status_code}: {response.text[:300]}") data = response.json() @@ -420,13 +436,8 @@ async def _acall_anthropic( async def _acall_ollama( self, messages: List[Dict], system: Optional[str], agent_name: str = "unknown" ) -> str: - url = f"{self.config.base_url or 'http://localhost:11434'}/api/chat" - payload = { - "model": self.config.model, - "messages": messages, - "stream": False, - } - async with httpx.AsyncClient(timeout=60) as client: + url, payload = self._ollama_request(messages, system) + async with httpx.AsyncClient(timeout=OLLAMA_TIMEOUT) as client: response = await client.post(url, json=payload) response.raise_for_status() data = response.json() diff --git a/cyberai/core/orchestrator.py b/cyberai/core/orchestrator.py index 858ce0b..4c44bbe 100644 --- a/cyberai/core/orchestrator.py +++ b/cyberai/core/orchestrator.py @@ -400,7 +400,9 @@ async def _dispatch_async(self, session: ScanSession, phase: ScanPhase) -> Dict[ async def _run_recon_async(self, session: ScanSession) -> Dict: from cyberai.agents.recon.async_agent import AsyncReconAgent + from cyberai.agents.recon.dns_tool import run_whois from cyberai.agents.recon.llm_detector import detect_llm_endpoints + from cyberai.agents.recon.subdomain_enum import fqdns from cyberai.core.types import OpenPort, ReconResult agent = AsyncReconAgent() @@ -418,17 +420,33 @@ async def _run_recon_async(self, session: ScanSession) -> Dict: llm_result = await asyncio.to_thread(detect_llm_endpoints, session.target) session.kb.set("recon.llm_endpoints", llm_result, agent="async_recon") + # whois too: AsyncReconAgent runs only nmap/dns/subdomains/tls, so the + # sync agent's recon.whois key and ReconResult.whois were empty on the + # async path. Blocking lookup -> offload, same as the LLM detector. + whois_result = await asyncio.to_thread(run_whois, session.target) + session.kb.set("recon.whois", whois_result, agent="async_recon") + + # HTTP attack surface: AsyncReconAgent has no web branch, so + # recon.web_surface was absent on the async path and every consumer + # of it (build_kb_graph, ExploitAgent) silently saw nothing. Reuse the + # sync agent's method rather than re-deriving the crawl and its + # findings here — one source of truth for the web surface. + if getattr(self.config, "use_web_recon", False): + from cyberai.agents.recon.agent import ReconAgent + + web_agent = ReconAgent(self.config, session, None, getattr(self, "audit", None)) + await asyncio.to_thread(web_agent._run_web_recon, session.target) + # Validated ReconResult so the planner KB graph gets port/service/ # subdomain nodes (build_kb_graph reads recon.result), matching sync. nmap = result.get("nmap") if isinstance(result.get("nmap"), dict) else {} raw_ports = nmap.get("ports", []) if isinstance(nmap, dict) else [] subs = result.get("subdomains") if isinstance(result.get("subdomains"), dict) else {} - subdomains = [ - r["fqdn"] for r in (subs.get("found") or []) if isinstance(r, dict) and r.get("fqdn") - ] + subdomains = fqdns(subs) recon_result = ReconResult( target=session.target, ports=[OpenPort(**p) for p in raw_ports if isinstance(p, dict)], + whois=whois_result if isinstance(whois_result, dict) else {}, dns=result.get("dns") if isinstance(result.get("dns"), dict) else {}, subdomains=subdomains, ) diff --git a/cyberai/mcp/tools.py b/cyberai/mcp/tools.py index 6603d44..fe66a5d 100644 --- a/cyberai/mcp/tools.py +++ b/cyberai/mcp/tools.py @@ -36,10 +36,10 @@ def register( # ── recon tools ─────────────────────────────────────────────────────── from cyberai.agents.recon.dns_tool import ( # noqa: E402 - detect_subdomains, run_dns, run_whois, ) +from cyberai.agents.recon.subdomain_enum import enumerate_subdomains # noqa: E402 from cyberai.agents.recon.nmap_tool import run_nmap # noqa: E402 register( @@ -88,8 +88,8 @@ def register( def _subdomain_handler(target: str, wordlist: list[str] | None = None) -> dict: - """Adapt detect_subdomains to MCP (wordlist optional).""" - return detect_subdomains(target, wordlist) + """Adapt enumerate_subdomains to MCP (wordlist optional).""" + return enumerate_subdomains(target, wordlist) register( diff --git a/tests/integration/test_async_pipeline.py b/tests/integration/test_async_pipeline.py index f08cc98..c99e8cb 100644 --- a/tests/integration/test_async_pipeline.py +++ b/tests/integration/test_async_pipeline.py @@ -255,6 +255,89 @@ def test_async_recon_populates_result_and_llm_for_planner_graph(self): ntypes = {d["ntype"] for _, d in graph.nodes(data=True)} assert "port" in ntypes and "llm_endpoint" in ntypes + def test_async_recon_writes_whois_like_sync(self): + """AsyncReconAgent runs only nmap/dns/subdomains/tls, so whois has to + be offloaded by the orchestrator or recon.whois and ReconResult.whois + stay empty on the async path while the sync agent fills both.""" + recon_payload = {"target": "t.local", "nmap": {"ports": []}, "dns": {}} + whois_payload = {"target": "t.local", "registrar": "Example Registrar"} + from cyberai.core.config import CyberAIConfig + from cyberai.core.orchestrator import AsyncOrchestrator + from cyberai.core.scan_session import ScanSession + + orch = AsyncOrchestrator(config=CyberAIConfig(), dry_run=False) + session = ScanSession(target="t.local") + with ( + patch( + "cyberai.agents.recon.async_agent.AsyncReconAgent.run", + new_callable=AsyncMock, + return_value=recon_payload, + ), + patch( + "cyberai.agents.recon.llm_detector.detect_llm_endpoints", + return_value={}, + ), + patch("cyberai.agents.recon.dns_tool.run_whois", return_value=whois_payload), + ): + asyncio.run(orch._run_recon_async(session)) + + assert session.kb.get("recon.whois") == whois_payload + assert session.kb.get("recon.result")["whois"] == whois_payload + + def test_async_recon_runs_web_surface_when_flag_is_on(self): + """AsyncReconAgent has no web branch. Without the orchestrator running + one, recon.web_surface is absent on the async path and build_kb_graph + plus ExploitAgent see an empty surface — the whole web layer is dead + under AsyncOrchestrator while it works under the sync one.""" + from cyberai.core.config import CyberAIConfig + from cyberai.core.orchestrator import AsyncOrchestrator + from cyberai.core.scan_session import ScanSession + + recon_payload = {"target": "t.local", "nmap": {"ports": []}, "dns": {}} + surface = { + "endpoints": [{"method": "GET", "url": "http://t.local/search", "params": ["q"]}] + } + orch = AsyncOrchestrator(config=CyberAIConfig(use_web_recon=True), dry_run=False) + session = ScanSession(target="t.local") + with ( + patch( + "cyberai.agents.recon.async_agent.AsyncReconAgent.run", + new_callable=AsyncMock, + return_value=recon_payload, + ), + patch("cyberai.agents.recon.llm_detector.detect_llm_endpoints", return_value={}), + patch("cyberai.agents.recon.dns_tool.run_whois", return_value={}), + patch("cyberai.agents.recon.agent.discover_surface", return_value=surface), + ): + asyncio.run(orch._run_recon_async(session)) + + assert session.kb.get("recon.web_surface") == surface + + def test_async_recon_skips_web_surface_when_flag_is_off(self): + """Default config must not start crawling: the flag gates the async + path exactly as it gates the sync one.""" + from cyberai.core.config import CyberAIConfig + from cyberai.core.orchestrator import AsyncOrchestrator + from cyberai.core.scan_session import ScanSession + + recon_payload = {"target": "t.local", "nmap": {"ports": []}, "dns": {}} + orch = AsyncOrchestrator(config=CyberAIConfig(), dry_run=False) + session = ScanSession(target="t.local") + with ( + patch( + "cyberai.agents.recon.async_agent.AsyncReconAgent.run", + new_callable=AsyncMock, + return_value=recon_payload, + ), + patch("cyberai.agents.recon.llm_detector.detect_llm_endpoints", return_value={}), + patch("cyberai.agents.recon.dns_tool.run_whois", return_value={}), + patch("cyberai.agents.recon.agent.discover_surface") as crawl, + ): + asyncio.run(orch._run_recon_async(session)) + + crawl.assert_not_called() + assert session.kb.get("recon.web_surface") is None + def test_sync_intel_runs_under_to_thread(self): """Intel/exploit/report stay sync; AsyncOrchestrator must offload them.""" from cyberai.core.scan_session import ScanSession, ScanPhase diff --git a/tests/integration/test_plan_web_order_pipeline.py b/tests/integration/test_plan_web_order_pipeline.py index fabef85..9d531e3 100644 --- a/tests/integration/test_plan_web_order_pipeline.py +++ b/tests/integration/test_plan_web_order_pipeline.py @@ -67,7 +67,7 @@ def send(url: str, method: str, params: Dict[str, str]) -> str: patch("cyberai.agents.recon.agent.run_nmap", return_value=RECON_RESULT), patch("cyberai.agents.recon.agent.run_whois", return_value={}), patch("cyberai.agents.recon.agent.run_dns", return_value={}), - patch("cyberai.agents.recon.agent.detect_subdomains", return_value={}), + patch("cyberai.agents.recon.agent.enumerate_subdomains", return_value={}), patch("cyberai.agents.recon.agent.detect_llm_endpoints", return_value={}), patch("cyberai.agents.recon.agent.discover_surface", return_value=SURFACE), patch("cyberai.agents.exploit.web_exploit._default_sender", return_value=send), diff --git a/tests/unit/test_behavioral_probe.py b/tests/unit/test_behavioral_probe.py index be45d2e..c286fbb 100644 --- a/tests/unit/test_behavioral_probe.py +++ b/tests/unit/test_behavioral_probe.py @@ -128,7 +128,7 @@ def _run_recon(config): patch("cyberai.agents.recon.agent.run_nmap", return_value=ok_nmap), patch("cyberai.agents.recon.agent.run_whois", return_value={}), patch("cyberai.agents.recon.agent.run_dns", return_value={}), - patch("cyberai.agents.recon.agent.detect_subdomains", return_value={}), + patch("cyberai.agents.recon.agent.enumerate_subdomains", return_value={}), patch("cyberai.agents.recon.agent.detect_llm_endpoints", return_value={}), patch( "cyberai.agents.recon.agent.build_probe_context", @@ -191,7 +191,7 @@ def fake_build(tgt, ports, mass_open=False): patch("cyberai.agents.recon.agent.run_nmap", return_value=mass_nmap), patch("cyberai.agents.recon.agent.run_whois", return_value={}), patch("cyberai.agents.recon.agent.run_dns", return_value={}), - patch("cyberai.agents.recon.agent.detect_subdomains", return_value={}), + patch("cyberai.agents.recon.agent.enumerate_subdomains", return_value={}), patch("cyberai.agents.recon.agent.detect_llm_endpoints", return_value={}), patch("cyberai.agents.recon.agent.build_probe_context", side_effect=fake_build), ): diff --git a/tests/unit/test_llm_client.py b/tests/unit/test_llm_client.py index 183fbbb..a441cbc 100644 --- a/tests/unit/test_llm_client.py +++ b/tests/unit/test_llm_client.py @@ -145,3 +145,61 @@ async def post(self, url, json): assert client.cost_tracker.call_count == 1 assert client.cost_tracker.total_input_tokens == 10 assert client.cost_tracker.total_output_tokens == 3 + + +def test_ollama_sync_and_async_send_identical_payloads(monkeypatch): + """Both ollama entry points must build the same request body. The async + one used to construct its own, dropping the system prompt and the raised + num_ctx, so the same prompt behaved differently depending on the path.""" + import asyncio + + from cyberai.core import llm_client as lc + + captured = {} + + class _Resp: + status_code = 200 + + def raise_for_status(self): + pass + + def json(self): + return {"message": {"content": "ok"}} + + def fake_post(url, json, timeout): + captured["sync"] = (url, json, timeout) + return _Resp() + + class _AsyncClient: + def __init__(self, timeout): + captured["async_timeout"] = timeout + + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + async def post(self, url, json): + captured["async"] = (url, json) + return _Resp() + + monkeypatch.setattr(lc.httpx, "post", fake_post) + monkeypatch.setattr(lc.httpx, "AsyncClient", _AsyncClient) + + client = lc.LLMClient.__new__(lc.LLMClient) + client.config = type("C", (), {"base_url": None, "model": "qwen2.5:7b", "provider": "ollama"})() + client.cost_tracker = None + client.budget_usd = 0 + + messages = [{"role": "user", "content": "hi"}] + client._call_ollama(messages, "SYS") + asyncio.run(client._acall_ollama(messages, "SYS")) + + sync_url, sync_payload, sync_timeout = captured["sync"] + async_url, async_payload = captured["async"] + assert sync_url == async_url + assert sync_payload == async_payload + assert sync_timeout == captured["async_timeout"] == lc.OLLAMA_TIMEOUT + assert async_payload["messages"][0] == {"role": "system", "content": "SYS"} + assert async_payload["options"]["num_ctx"] == 8192 diff --git a/tests/unit/test_recon.py b/tests/unit/test_recon.py index 982a312..df067c4 100644 --- a/tests/unit/test_recon.py +++ b/tests/unit/test_recon.py @@ -52,7 +52,7 @@ def test_recon_logs_nmap_failure_visibly(monkeypatch): patch("cyberai.agents.recon.agent.run_nmap", return_value=err_nmap), patch("cyberai.agents.recon.agent.run_whois", return_value={}), patch("cyberai.agents.recon.agent.run_dns", return_value={}), - patch("cyberai.agents.recon.agent.detect_subdomains", return_value={}), + patch("cyberai.agents.recon.agent.enumerate_subdomains", return_value={}), patch("cyberai.agents.recon.agent.detect_llm_endpoints", return_value={}), ): agent.run("t.local") @@ -84,10 +84,52 @@ def test_recon_logs_nmap_success(monkeypatch): patch("cyberai.agents.recon.agent.run_nmap", return_value=ok_nmap), patch("cyberai.agents.recon.agent.run_whois", return_value={}), patch("cyberai.agents.recon.agent.run_dns", return_value={}), - patch("cyberai.agents.recon.agent.detect_subdomains", return_value={}), + patch("cyberai.agents.recon.agent.enumerate_subdomains", return_value={}), patch("cyberai.agents.recon.agent.detect_llm_endpoints", return_value={}), ): agent.run("t.local") assert "nmap_scan complete" in logged assert not any("nmap_scan FAILED" in m for m in logged) + + +class TestSubdomainShapeParity: + """The sync ReconAgent and AsyncOrchestrator must write the same shape + under recon.subdomains and derive ReconResult.subdomains identically. + They used to call different enumerators returning different keys, so a + consumer of the KB key was only ever correct for one of the two paths.""" + + _ENUM_RESULT = { + "domain": "t.local", + "found": [{"fqdn": "api.t.local", "ips": ["1.2.3.4"]}], + "count": 1, + "checked": 1, + "wordlist": ["api"], + } + + def test_sync_recon_writes_found_shape_and_fills_recon_result(self): + from unittest.mock import MagicMock, patch + + from cyberai.agents.recon.agent import ReconAgent + from cyberai.core.config import CyberAIConfig + from cyberai.core.scan_session import ScanSession + + session = ScanSession(target="t.local") + agent = ReconAgent(CyberAIConfig(), session, MagicMock(), MagicMock()) + with ( + patch( + "cyberai.agents.recon.agent.run_nmap", + return_value={"target": "t.local", "ports": []}, + ), + patch("cyberai.agents.recon.agent.run_whois", return_value={}), + patch("cyberai.agents.recon.agent.run_dns", return_value={}), + patch( + "cyberai.agents.recon.agent.enumerate_subdomains", + return_value=self._ENUM_RESULT, + ), + patch("cyberai.agents.recon.agent.detect_llm_endpoints", return_value={}), + ): + agent.run("t.local") + + assert session.kb.get("recon.subdomains") == self._ENUM_RESULT + assert session.kb.get("recon.result")["subdomains"] == ["api.t.local"] diff --git a/tests/unit/test_recon_guards.py b/tests/unit/test_recon_guards.py index ee2130f..5a37ea3 100644 --- a/tests/unit/test_recon_guards.py +++ b/tests/unit/test_recon_guards.py @@ -14,7 +14,7 @@ _RECON_TOOLS = { "run_whois": {}, "run_dns": {}, - "detect_subdomains": {}, + "enumerate_subdomains": {}, "detect_llm_endpoints": {}, } diff --git a/tests/unit/test_subdomain_enum.py b/tests/unit/test_subdomain_enum.py index 9e73734..f4cae77 100644 --- a/tests/unit/test_subdomain_enum.py +++ b/tests/unit/test_subdomain_enum.py @@ -1,6 +1,7 @@ from unittest.mock import patch from cyberai.agents.recon.subdomain_enum import ( enumerate_subdomains, + fqdns, _resolve, load_wordlist, DEFAULT_WORDLIST, @@ -81,3 +82,34 @@ def test_load_wordlist_from_file(tmp_path): wl.write_text("sub1\nsub2\n# comment\nsub3\n") result = load_wordlist(str(wl)) assert result == ["sub1", "sub2", "sub3"] + + +class TestFqdnsContract: + """fqdns() is the single place that turns an enumerator result into bare + hostnames. Both the sync ReconAgent and AsyncOrchestrator feed + ReconResult.subdomains through it, so a shape change cannot make one + pipeline path silently disagree with the other.""" + + def test_extracts_hostnames_from_found(self): + result = {"found": [{"fqdn": "api.t.local"}, {"fqdn": "dev.t.local"}]} + assert fqdns(result) == ["api.t.local", "dev.t.local"] + + def test_skips_malformed_entries(self): + result = {"found": [{"fqdn": "api.t.local"}, {"ips": ["1.2.3.4"]}, "www.t.local", {}]} + assert fqdns(result) == ["api.t.local"] + + def test_empty_inputs_yield_empty_list(self): + assert fqdns(None) == [] + assert fqdns({}) == [] + assert fqdns({"found": []}) == [] + + def test_real_enumerator_output_is_consumable(self): + """Guards the contract against the producer, not a hand-built dict.""" + from unittest.mock import patch + + with patch( + "cyberai.agents.recon.subdomain_enum._resolve", + side_effect=lambda fqdn: {"fqdn": fqdn, "ips": ["1.2.3.4"]}, + ): + produced = enumerate_subdomains("t.local", wordlist=["api", "dev"]) + assert sorted(fqdns(produced)) == ["api.t.local", "dev.t.local"] diff --git a/tests/unit/test_web_wiring.py b/tests/unit/test_web_wiring.py index cbfd5bc..1d436de 100644 --- a/tests/unit/test_web_wiring.py +++ b/tests/unit/test_web_wiring.py @@ -34,7 +34,7 @@ "run_nmap": {"target": "t.local", "ports": []}, "run_whois": {}, "run_dns": {}, - "detect_subdomains": {}, + "enumerate_subdomains": {}, "detect_llm_endpoints": {}, } @@ -51,7 +51,7 @@ def recon_patches(): patch("cyberai.agents.recon.agent.run_nmap", return_value=_RECON_PATCHES["run_nmap"]), patch("cyberai.agents.recon.agent.run_whois", return_value={}), patch("cyberai.agents.recon.agent.run_dns", return_value={}), - patch("cyberai.agents.recon.agent.detect_subdomains", return_value={}), + patch("cyberai.agents.recon.agent.enumerate_subdomains", return_value={}), patch("cyberai.agents.recon.agent.detect_llm_endpoints", return_value={}), ): yield