diff --git a/cyberai/agents/redteam/__init__.py b/cyberai/agents/redteam/__init__.py index 3b12b4c..5c25e71 100644 --- a/cyberai/agents/redteam/__init__.py +++ b/cyberai/agents/redteam/__init__.py @@ -1,5 +1,6 @@ """LLM offensive red-team fuzzing package.""" +from .agent import RedTeamAgent from .fuzzer import FuzzReport, FuzzResult, LLMChannelFuzzer, SendFn from .payloads import ( ACK_PREFIX, @@ -17,6 +18,7 @@ "InjectionPayload", "LLMChannelFuzzer", "PayloadCategory", + "RedTeamAgent", "SendFn", "build_corpus", "full_corpus", diff --git a/cyberai/agents/redteam/agent.py b/cyberai/agents/redteam/agent.py new file mode 100644 index 0000000..11ca993 --- /dev/null +++ b/cyberai/agents/redteam/agent.py @@ -0,0 +1,111 @@ +"""Red-team agent: drive the injection corpus at planned LLM channels. + +The fuzzer is transport-agnostic — it maps a payload string to a response +string and knows nothing about HTTP. This agent supplies the missing half: it +reads the LLM/RAG endpoints the planner named, builds an HTTP channel for +each, and records what the fuzzer confirms as findings on the session. + +Confirmation tiers come straight from the fuzzer and are not re-interpreted +here: an out-of-band callback is the only thing that earns full confidence, +and an echoed marker or a leak heuristic stays below it. Promoting a +heuristic to a confirmed finding is exactly the false positive the OOB path +exists to avoid. +""" + +from __future__ import annotations + +from typing import Any, Callable, Dict, List, Optional + +import httpx + +from cyberai.core.base_agent import BaseAgent + +from .fuzzer import FuzzReport, LLMChannelFuzzer + +DEFAULT_TIMEOUT = 10.0 + +# Field names chat and RAG APIs commonly accept for the user's turn. The +# payload is sent under each in one body: a wrong key is ignored, a right one +# is read, and one request beats probing the schema. +_PROMPT_FIELDS = ("prompt", "message", "input", "query", "text") + + +def _default_channel(url: str, timeout: float = DEFAULT_TIMEOUT) -> Callable[[str], str]: + """Return a send function that POSTs a payload to `url` as JSON.""" + + def _send(payload: str) -> str: + body = {field: payload for field in _PROMPT_FIELDS} + try: + with httpx.Client(timeout=timeout, follow_redirects=True) as client: + return client.post(url, json=body).text + except Exception: + return "" + + return _send + + +class RedTeamAgent(BaseAgent): + """Fuzz the LLM channels a plan names; record confirmed injections.""" + + AGENT_NAME = "redteam" + ROLE = "LLM Red Team" + + def _register_tools(self) -> None: # channels are driven directly + pass + + def run(self, target: str, context: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: + self._check_iteration_limit() + urls = self._planned_channels() + if not urls: + self._log("No injection-fuzz subtasks in plan — red team skipped") + return {"channels": 0, "confirmed": 0, "reports": []} + + fuzzer = (context or {}).get("fuzzer") or LLMChannelFuzzer() + channel_factory = (context or {}).get("channel_factory") or _default_channel + + reports: List[Dict[str, Any]] = [] + confirmed = 0 + for url in urls: + report = fuzzer.fuzz_channel(channel_factory(url), channel_id=url) + confirmed += report.confirmed_count + self._record(report, url) + reports.append(report.to_dict()) + + self.kb.set("redteam.reports", reports, agent=self.AGENT_NAME) + self._log(f"fuzzed {len(urls)} channel(s), {confirmed} confirmed") + return {"channels": len(urls), "confirmed": confirmed, "reports": reports} + + def _planned_channels(self) -> List[str]: + """URLs from injection-fuzz subtasks, in plan order, deduplicated.""" + plan = self.kb.get("plan") or {} + todo = plan.get("todo") if isinstance(plan, dict) else None + if not isinstance(todo, list): + return [] + urls: List[str] = [] + for task in todo: + if not isinstance(task, dict) or task.get("action") != "injection-fuzz": + continue + url = task.get("target") + if url and url not in urls: + urls.append(str(url)) + return urls + + def _record(self, report: FuzzReport, url: str) -> None: + """Turn fuzz results into session findings, one per signal.""" + from cyberai.core.scan_session import Severity + + for result in report.results: + if result.severity == "INFO": + continue + finding = self.session.add_finding( + severity=Severity(result.severity), + title=f"Prompt injection ({result.category}) on LLM endpoint", + description=(f"Payload '{result.payload_id}' delivered to {url}: {result.detail}."), + agent=self.AGENT_NAME, + target=url, + evidence=[result.to_dict()], + data=result.to_dict(), + ) + # Only an out-of-band callback proves the injection executed; + # everything else is a signal in the response text. + finding.confidence = 1.0 if result.oob_confirmed else 0.6 diff --git a/cyberai/core/config.py b/cyberai/core/config.py index 2d9388f..678f4a2 100644 --- a/cyberai/core/config.py +++ b/cyberai/core/config.py @@ -122,6 +122,8 @@ class CyberAIConfig: use_api_discovery: bool = False # Flag-gated: attack the endpoints the planner named before the rest. use_plan_web_order: bool = False + # Flag-gated: fuzz the LLM channels the planner named, in the exploit phase. + use_planned_redteam: bool = False exploit_memory_path: Optional[str] = None # Flag-gated: parse local practice-lab machine artifacts and detect flags. use_lab_dogfood: bool = False @@ -184,6 +186,7 @@ def from_env(cls) -> "CyberAIConfig": use_web_exploit=_env_bool("CYBERAI_USE_WEB_EXPLOIT", False), use_api_discovery=_env_bool("CYBERAI_USE_API_DISCOVERY", False), use_plan_web_order=_env_bool("CYBERAI_USE_PLAN_WEB_ORDER", False), + use_planned_redteam=_env_bool("CYBERAI_USE_PLANNED_REDTEAM", False), use_lab_dogfood=_env_bool("CYBERAI_USE_LAB_DOGFOOD", False), web_enable_bench_trigger=_env_bool("CYBERAI_WEB_ENABLE_BENCH_TRIGGER", False), air_gapped=_env_bool("CYBERAI_AIR_GAPPED", False), diff --git a/cyberai/core/orchestrator.py b/cyberai/core/orchestrator.py index 391a8eb..858ce0b 100644 --- a/cyberai/core/orchestrator.py +++ b/cyberai/core/orchestrator.py @@ -279,9 +279,26 @@ def _run_exploit(self, session: ScanSession) -> Dict: agent = ExploitAgent(self.config, session, self._client_for(ScanPhase.EXPLOIT), self.audit) result = agent.run(session.target) + redteam = self._run_planned_redteam(session) + if redteam: + result = {**result, "redteam": redteam} session.kb_set("exploit", result) return result + def _run_planned_redteam(self, session: ScanSession) -> Dict: + """Fuzz the LLM channels the plan named, inside the exploit phase. + + Not a phase of its own: the planner already decides whether there is + anything to fuzz, and a phase that exists only to be skipped costs the + CLI, the replay format, and the scorecard a column each. + """ + if not getattr(self.config, "use_planned_redteam", False): + return {} + from cyberai.agents.redteam.agent import RedTeamAgent + + agent = RedTeamAgent(self.config, session, self._client_for(ScanPhase.EXPLOIT), self.audit) + return agent.run(session.target) + def _run_report(self, session: ScanSession) -> Dict: from cyberai.agents.report.agent import ReportAgent from cyberai.agents.report.html_renderer import render_html_report diff --git a/tests/integration/test_plan_web_order_pipeline.py b/tests/integration/test_plan_web_order_pipeline.py new file mode 100644 index 0000000..fabef85 --- /dev/null +++ b/tests/integration/test_plan_web_order_pipeline.py @@ -0,0 +1,97 @@ +"""The planner's endpoint order survives the whole pipeline. + +Each link was proven in isolation: the graph gains HTTP endpoint nodes, the +planner emits web-exploit subtasks, `exploit_surface` accepts a priority, and +the exploit agent reads the plan. None of that proves they meet. A phase +reordering, a KB key rename, or a plan written after the exploit phase reads +it would leave every unit test green and the order silently dropped, which is +exactly the failure this pipeline was built to avoid. +""" + +from __future__ import annotations + +from typing import Dict, List +from unittest.mock import MagicMock, patch + +from cyberai.core.config import CyberAIConfig +from cyberai.core.orchestrator import Orchestrator +from cyberai.core.scan_session import ScanPhase + +TARGET = "t.local" +BASE = "http://t.local" +# `dull` carries no hinted parameter name, so the deterministic order would +# put it last; naming it first in the plan is what has to survive. +DULL = f"{BASE}/dull" +HINTED = f"{BASE}/read" + +SURFACE = { + "base_url": BASE, + "reachable": True, + "pages_fetched": 1, + # Order matters: the plan follows graph insertion order, while the + # deterministic ranking follows the parameter name hint. Listing the dull + # endpoint first makes the two disagree, so the attack order proves which + # one actually ran instead of agreeing with both. + "endpoints": [ + {"url": DULL, "method": "POST", "params": ["zzz"], "source": "hint"}, + {"url": HINTED, "method": "GET", "params": ["path"], "source": "hint"}, + ], + "routes": [], + "spec_url": None, +} + +RECON_RESULT = {"target": TARGET, "ports": [], "subdomains": []} + + +def _config(**flags) -> CyberAIConfig: + return CyberAIConfig( + enable_planner=True, + use_web_exploit=True, + **flags, + ) + + +def _run(config: CyberAIConfig) -> List[str]: + """Run recon->plan->exploit with a stand-in sender; return attack order.""" + hit: List[str] = [] + + def send(url: str, method: str, params: Dict[str, str]) -> str: + hit.append(url) + return "" + + orch = Orchestrator( + config, + phases=[ScanPhase.RECON, ScanPhase.EXPLOIT], + ) + with ( + 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.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), + patch("cyberai.agents.exploit.web_exploit._default_json_sender", return_value=send), + patch("cyberai.core.orchestrator.Orchestrator._client_for", return_value=MagicMock()), + ): + session = orch.run(TARGET) + + return list(dict.fromkeys(hit)), session + + +def test_plan_phase_is_inserted_and_writes_web_subtasks(): + _, session = _run(_config(use_web_recon=True, use_plan_web_order=True)) + todo = (session.kb.get("plan") or {}).get("todo") or [] + web = [(t["target"], t["method"]) for t in todo if t["action"] == "web-exploit"] + assert web == [(DULL, "POST"), (HINTED, "GET")] + + +def test_plan_order_reaches_the_attack(): + order, _ = _run(_config(use_web_recon=True, use_plan_web_order=True)) + # Against the name hint, which would put HINTED first. + assert order == [DULL, HINTED] + + +def test_without_the_flag_the_deterministic_order_runs(): + order, _ = _run(_config(use_web_recon=True)) + assert order == [HINTED, DULL] diff --git a/tests/unit/test_redteam_agent.py b/tests/unit/test_redteam_agent.py new file mode 100644 index 0000000..fc3469a --- /dev/null +++ b/tests/unit/test_redteam_agent.py @@ -0,0 +1,194 @@ +"""The red-team agent drives planned LLM channels and records what holds.""" + +from __future__ import annotations + +from typing import Any, Dict, List +from unittest.mock import MagicMock + +from cyberai.agents.redteam.agent import RedTeamAgent, _default_channel +from cyberai.agents.redteam.fuzzer import FuzzReport, FuzzResult +from cyberai.core.config import CyberAIConfig +from cyberai.core.scan_session import ScanSession, Severity + +CHAT = "http://t.local/chat" +RAG = "http://t.local/v1/completions" + +PLAN = { + "todo": [ + {"action": "exploit", "target": "CVE-1"}, + {"action": "injection-fuzz", "target": CHAT}, + {"action": "injection-fuzz", "target": RAG}, + ] +} + + +class _StubFuzzer: + """Stand-in fuzzer: records channels and replays canned results.""" + + def __init__(self, results: List[FuzzResult] | None = None): + self.channels: List[str] = [] + self.sent: List[str] = [] + self._results = results or [] + + def fuzz_channel(self, send_fn, channel_id: str = "") -> FuzzReport: + self.channels.append(channel_id) + self.sent.append(send_fn("payload-probe")) + return FuzzReport(channel_id=channel_id, results=list(self._results)) + + +def _agent(plan: Any = None) -> RedTeamAgent: + session = ScanSession(target="t.local") + if plan is not None: + session.kb_set("plan", plan) + return RedTeamAgent(CyberAIConfig(), session, MagicMock(), MagicMock()) + + +def _run(agent: RedTeamAgent, fuzzer: Any, **ctx: Any) -> Dict[str, Any]: + return agent.run("t.local", {"fuzzer": fuzzer, "channel_factory": lambda u: lambda p: u, **ctx}) + + +def test_channels_follow_plan_order(): + agent = _agent(PLAN) + fuzzer = _StubFuzzer() + result = _run(agent, fuzzer) + assert fuzzer.channels == [CHAT, RAG] + assert result["channels"] == 2 + + +def test_duplicate_targets_are_fuzzed_once(): + plan = {"todo": [{"action": "injection-fuzz", "target": CHAT}] * 3} + fuzzer = _StubFuzzer() + _run(_agent(plan), fuzzer) + assert fuzzer.channels == [CHAT] + + +def test_no_plan_or_no_fuzz_subtasks_does_nothing(): + for plan in (None, "not-a-dict", {"todo": []}, {"todo": [{"action": "exploit"}]}): + fuzzer = _StubFuzzer() + result = _run(_agent(plan), fuzzer) + assert fuzzer.channels == [] + assert result == {"channels": 0, "confirmed": 0, "reports": []} + + +def test_confirmed_result_becomes_a_finding_with_full_confidence(): + agent = _agent(PLAN) + oob = FuzzResult( + payload_id="p1", + category="exfiltration", + oob_confirmed=True, + severity="CRITICAL", + detail="out-of-band callback confirmed", + ) + result = _run(agent, _StubFuzzer([oob])) + + assert result["confirmed"] == 2 # one per channel + findings = agent.session.findings + assert len(findings) == 2 + assert findings[0].severity is Severity.CRITICAL + assert findings[0].target == CHAT + assert findings[0].confidence == 1.0 + + +def test_unconfirmed_signal_is_recorded_below_full_confidence(): + agent = _agent({"todo": [{"action": "injection-fuzz", "target": CHAT}]}) + echoed = FuzzResult(payload_id="p2", category="jailbreak", ack_echoed=True, severity="HIGH") + _run(agent, _StubFuzzer([echoed])) + + finding = agent.session.findings[0] + assert finding.severity is Severity.HIGH + assert finding.confidence < 1.0 + + +def test_info_results_produce_no_finding(): + agent = _agent({"todo": [{"action": "injection-fuzz", "target": CHAT}]}) + quiet = FuzzResult(payload_id="p3", category="jailbreak", severity="INFO") + _run(agent, _StubFuzzer([quiet])) + assert agent.session.findings == [] + + +def test_reports_are_stored_in_the_knowledge_base(): + agent = _agent(PLAN) + _run(agent, _StubFuzzer()) + stored = agent.session.kb.get("redteam.reports") + assert [r["channel_id"] for r in stored] == [CHAT, RAG] + + +def test_default_channel_posts_the_payload_under_known_field_names(): + sent: Dict[str, Any] = {} + + class _Resp: + text = "answer" + + class _Client: + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def post(self, url, json): + sent["url"] = url + sent["json"] = json + return _Resp() + + import cyberai.agents.redteam.agent as mod + + original = mod.httpx.Client + mod.httpx.Client = lambda **kw: _Client() + try: + assert _default_channel(CHAT)("INJECT") == "answer" + finally: + mod.httpx.Client = original + + assert sent["url"] == CHAT + assert set(sent["json"].values()) == {"INJECT"} + assert "prompt" in sent["json"] + + +def test_default_channel_survives_a_dead_target(): + import cyberai.agents.redteam.agent as mod + + original = mod.httpx.Client + + def _boom(**kw): + raise RuntimeError("connection refused") + + mod.httpx.Client = _boom + try: + assert _default_channel(CHAT)("INJECT") == "" + finally: + mod.httpx.Client = original + + +# ── orchestrator wiring ─────────────────────────────────────────────── + + +def test_orchestrator_skips_red_team_without_the_flag(): + from cyberai.core.orchestrator import Orchestrator + + orch = Orchestrator(CyberAIConfig()) + session = ScanSession(target="t.local") + session.kb_set("plan", PLAN) + assert orch._run_planned_redteam(session) == {} + + +def test_orchestrator_runs_red_team_when_enabled(): + from unittest.mock import patch + + from cyberai.core.orchestrator import Orchestrator + + orch = Orchestrator(CyberAIConfig(use_planned_redteam=True)) + orch.audit = MagicMock() + session = ScanSession(target="t.local") + session.kb_set("plan", PLAN) + + with ( + patch.object(Orchestrator, "_client_for", return_value=MagicMock()), + patch("cyberai.agents.redteam.agent.LLMChannelFuzzer", return_value=_StubFuzzer()), + # Without this the stub calls the real channel, which spends a full + # connect timeout per host proving nothing about the wiring. + patch("cyberai.agents.redteam.agent._default_channel", lambda url: lambda p: ""), + ): + result = orch._run_planned_redteam(session) + + assert result["channels"] == 2