diff --git a/cyberai/__main__.py b/cyberai/__main__.py index d3a30c0..35d73e9 100644 --- a/cyberai/__main__.py +++ b/cyberai/__main__.py @@ -1,8 +1,11 @@ +"""CyberAI CLI — entry point for the pentest pipeline.""" +from __future__ import annotations + import click from rich.console import Console from rich.panel import Panel + from .core.config import CyberAIConfig -from .core.session import PentestSession from .core.orchestrator import Orchestrator console = Console() @@ -19,45 +22,56 @@ [dim]AI-native pentest platform[/dim] """ + @click.group() -def cli(): - """CyberAI — AI-powered pentest platform""" - pass +def cli() -> None: + """CyberAI — AI-powered pentest platform.""" + @cli.command() @click.argument("target") @click.option("--verbose", "-v", is_flag=True, help="Verbose output") -@click.option("--provider", default="openai", help="LLM provider") -def scan(target: str, verbose: bool, provider: str): - """Run full pentest pipeline against TARGET""" +@click.option("--provider", default=None, help="LLM provider (openai/anthropic/ollama)") +@click.option("--dry-run", is_flag=True, help="Run pipeline without real network calls") +@click.option("--scope", multiple=True, help="Authorized scope entry (repeatable)") +def scan( + target: str, + verbose: bool, + provider: str | None, + dry_run: bool, + scope: tuple[str, ...], +) -> None: + """Run full pentest pipeline against TARGET.""" console.print(BANNER) console.print(Panel(f"[bold]Target:[/bold] {target}", style="red")) config = CyberAIConfig.from_env() config.verbose = verbose - config.llm.provider = provider + if provider: + config.llm.provider = provider - session = PentestSession(target=target) - orchestrator = Orchestrator(config) + orchestrator = Orchestrator(config=config, dry_run=dry_run) console.print("[yellow]→[/yellow] Starting pipeline...") - session = orchestrator.run_pipeline(session) + session = orchestrator.run(target, authorized_scope=list(scope)) console.print(f"\n[green]✓[/green] Done. Findings: {len(session.findings)}") summary = session.summary() - for k, v in summary.items(): - console.print(f" {k}: {v}") + for key, value in summary.items(): + console.print(f" {key}: {value}") + @cli.command() -def status(): - """Show CyberAI status and config""" +def status() -> None: + """Show CyberAI status and config.""" config = CyberAIConfig.from_env() console.print(Panel( f"Provider: {config.llm.provider}\n" f"Model: {config.llm.model}\n" f"Output: {config.output_dir}", - title="CyberAI Status" + title="CyberAI Status", )) + if __name__ == "__main__": cli() diff --git a/cyberai/core/orchestrator.py b/cyberai/core/orchestrator.py index 85f9ca6..4e73168 100644 --- a/cyberai/core/orchestrator.py +++ b/cyberai/core/orchestrator.py @@ -1,15 +1,24 @@ """ Orchestrator — coordinates the full multi-agent pipeline. ReconAgent → IntelAgent → ExploitAgent → ReportAgent + +Day 5 of STANDOFF rewrite: closes KI-1. + +The orchestrator now takes a CyberAIConfig, builds the shared LLMClient +and AuditLogger, and constructs every agent with the new BaseAgent +contract: Agent(config, session, llm, audit). """ from __future__ import annotations -from typing import Any, Dict, List + +from datetime import datetime, timezone +from typing import Any, Dict, List, Optional from rich.console import Console from rich.panel import Panel -from cyberai.core.scan_session import ScanSession, ScanPhase -from cyberai.core.logger import get_logger +from cyberai.core.config import CyberAIConfig +from cyberai.core.logger import AuditLogger, get_logger +from cyberai.core.scan_session import ScanPhase, ScanSession console = Console() log = get_logger("orchestrator") @@ -18,7 +27,7 @@ class Orchestrator: """ Runs the full CyberAI pipeline for a given target. - Phases are configurable — skip any by omitting from phases list. + Phases are configurable — skip any by omitting from the phases list. """ DEFAULT_PHASES = [ @@ -30,29 +39,48 @@ class Orchestrator: def __init__( self, - phases: List[ScanPhase] = None, - authorized_scope: List[str] = None, + config: Optional[CyberAIConfig] = None, + phases: Optional[List[ScanPhase]] = None, dry_run: bool = False, - ): - self.phases = phases or self.DEFAULT_PHASES - self.authorized_scope = authorized_scope or [] - self.dry_run = dry_run + ) -> None: + self.config = config or CyberAIConfig() + self.phases = phases or self.DEFAULT_PHASES + self.dry_run = dry_run + + # Shared LLM client — built lazily so dry-run never needs an API key. + self._llm = None + + # ── llm (lazy) ──────────────────────────────────────────────────── + + @property + def llm(self): + """Lazily build the shared LLMClient. Skipped entirely in dry-run.""" + if self._llm is None and not self.dry_run: + from cyberai.core.llm_client import LLMClient + self._llm = LLMClient(self.config.llm) + return self._llm - def run(self, target: str) -> ScanSession: + # ── public API ──────────────────────────────────────────────────── + + def run( + self, + target: str, + authorized_scope: Optional[List[str]] = None, + ) -> ScanSession: """ - Execute full pipeline for target. - Returns completed ScanSession with all results in KB. + Execute the full pipeline for `target`. + Returns the completed ScanSession with all results in its KB. """ session = ScanSession( target=target, - authorized_scope=self.authorized_scope, + authorized_scope=authorized_scope or [], ) console.print(Panel( f"[bold red]CyberAI Orchestrator[/bold red]\n" f"Target : [yellow]{target}[/yellow]\n" f"Phases : [yellow]{[p.value for p in self.phases]}[/yellow]\n" - f"Scope : [yellow]{self.authorized_scope or 'not set'}[/yellow]\n" + f"Scope : [yellow]{session.authorized_scope or 'not set'}[/yellow]\n" f"Dry Run : [yellow]{self.dry_run}[/yellow]\n" f"Session : [dim]{session.session_id}[/dim]", border_style="red", @@ -61,12 +89,14 @@ def run(self, target: str) -> ScanSession: session.start() log.info(f"Pipeline started — target={target} session={session.session_id}") + self.audit = AuditLogger(session_id=session.session_id) + for phase in self.phases: self._run_phase(session, phase) - if not session.phases[-1].success: + if session.phases and not session.phases[-1].success: log.warning(f"Phase {phase.value} failed — continuing pipeline") - if all(p.success for p in session.phases): + if session.phases and all(p.success for p in session.phases): session.complete() console.print("[bold green]✓ Pipeline complete[/bold green]") else: @@ -77,6 +107,8 @@ def run(self, target: str) -> ScanSession: log.info(f"Pipeline done — state={session.state.value}") return session + # ── phase execution ─────────────────────────────────────────────── + def _run_phase(self, session: ScanSession, phase: ScanPhase) -> None: console.print(f"\n[bold red]▶ {phase.value.upper()}[/bold red]") started = _now() @@ -91,7 +123,7 @@ def _run_phase(self, session: ScanSession, phase: ScanPhase) -> None: session.record_phase(phase, success=True, started=started, data=data) console.print(f"[green]✓ {phase.value} done[/green]") - except Exception as exc: + except Exception as exc: # noqa: BLE001 — pipeline must survive one bad phase session.record_phase( phase, success=False, started=started, error=str(exc) ) @@ -99,25 +131,28 @@ def _run_phase(self, session: ScanSession, phase: ScanPhase) -> None: log.error(f"Phase {phase.value} raised", exc_info=True) def _dispatch(self, session: ScanSession, phase: ScanPhase) -> Dict[str, Any]: - if phase == ScanPhase.RECON: - return self._run_recon(session) - if phase == ScanPhase.INTEL: - return self._run_intel(session) - if phase == ScanPhase.EXPLOIT: - return self._run_exploit(session) - if phase == ScanPhase.REPORT: - return self._run_report(session) - return {} + dispatch = { + ScanPhase.RECON: self._run_recon, + ScanPhase.INTEL: self._run_intel, + ScanPhase.EXPLOIT: self._run_exploit, + ScanPhase.REPORT: self._run_report, + } + handler = dispatch.get(phase) + return handler(session) if handler else {} + + # ── per-phase handlers ──────────────────────────────────────────── def _run_recon(self, session: ScanSession) -> Dict: from cyberai.agents.recon.agent import ReconAgent - result = ReconAgent(kb=session.kb).run(session.target) + agent = ReconAgent(self.config, session, self.llm, self.audit) + result = agent.run(session.target) session.kb_set("recon", result) return result def _run_intel(self, session: ScanSession) -> Dict: from cyberai.agents.intel.agent import IntelAgent - result = IntelAgent(kb=session.kb).run(session.target) + agent = IntelAgent(self.config, session, self.llm, self.audit) + result = agent.run(session.target) session.kb_set("intel", result) return result @@ -133,7 +168,8 @@ def _run_exploit(self, session: ScanSession) -> Dict: for w in v.warnings: console.print(f"[yellow]⚠ {w}[/yellow]") - result = ExploitAgent(kb=session.kb).run(session.target) + agent = ExploitAgent(self.config, session, self.llm, self.audit) + result = agent.run(session.target) session.kb_set("exploit", result) return result @@ -141,7 +177,8 @@ def _run_report(self, session: ScanSession) -> Dict: from cyberai.agents.report.agent import ReportAgent from cyberai.agents.report.html_renderer import render_html_report - result = ReportAgent(kb=session.kb).run(session.target) + agent = ReportAgent(self.config, session, self.llm, self.audit) + result = agent.run(session.target) session.kb_set("report", result) output = f"report_{session.session_id}.html" @@ -151,5 +188,4 @@ def _run_report(self, session: ScanSession) -> Dict: def _now() -> str: - from datetime import datetime, timezone return datetime.now(timezone.utc).isoformat() diff --git a/docs/architecture/known-issues.md b/docs/architecture/known-issues.md index 7df2d24..843157d 100644 --- a/docs/architecture/known-issues.md +++ b/docs/architecture/known-issues.md @@ -5,38 +5,40 @@ rewrite. Each item is fixed by a specific day; see `STANDOFF.md`. ## The Issues -### 🔴 KI-1 — CLI ↔ Orchestrator API mismatch -`__main__.py` calls `Orchestrator(config)` and `run_pipeline(session)` — -neither matches the actual API. **Fixed by:** Day 5. +### 🟢 KI-1 — CLI ↔ Orchestrator API mismatch ✅ FIXED IN DAY 5 +`Orchestrator` now takes `(config, phases, dry_run)`; `run(target, +authorized_scope)` owns session creation and builds the shared +`LLMClient`/`AuditLogger`. `__main__.py` calls the real API and gained +`--dry-run` / `--scope`. `python -m cyberai scan --dry-run` runs all +four phases and exits cleanly. Verified by +`tests/unit/test_orchestrator_config.py`. ### 🟢 KI-2 — Two competing session classes ✅ FIXED IN DAY 3 -`scan_session.py` is now the single source of truth. `session.py` is a -backward-compat shim. Verified by `tests/unit/test_session_shim.py`. ### 🟢 KI-3 — BaseAgent didn't match what agents use ✅ FIXED IN DAY 4 -`BaseAgent.__init__` now takes `(config, session, llm, audit)` and -exposes `self.session`, `self.kb`, `self.llm`, `self.memory`. Agents are -migrated to actually use this contract in day 6. Verified by -`tests/unit/test_base_agent.py`. ### 🟢 KI-4 — Agents called non-existent methods ✅ FIXED IN DAY 4 -`_check_iteration_limit()` and `_log()` now exist on `BaseAgent`. -`AgentMemory` (with `add()`/`to_messages()`) backs `self.memory`. -`self.llm.chat()` is addressed in day 6 when ExploitAgent is migrated to -`self.llm.call()`. Verified by `tests/unit/test_base_agent.py`. +`_check_iteration_limit()`, `_log()`, `AgentMemory` exist on +`BaseAgent`. `self.llm.chat()` remains — addressed in day 6 when agents +are migrated to `self.llm.call()`. -### 🔴 KI-5 — Finding signature mismatch ✅ FIXED IN DAY 3 +### 🟢 KI-5 — Finding signature mismatch ✅ FIXED IN DAY 3 ### 🟢 KI-6 — Tool param name mismatch ✅ FIXED IN DAY 4 -`Tool` accepts both `params` and `parameters`, synced via -`__post_init__`. All agents register tools with `parameters=...` so this -closed without touching any agent file. ### 🔴 KI-7 — `LLMClient.chat()` doesn't exist -Actual method is `call()`. **Fixed by:** Day 6. +`ExploitAgent` calls `self.llm.chat()`; the real method is `call()`. +Agents still use the old `BaseAgent` construction internally — they are +migrated to the new contract in day 6. **Fixed by:** Day 6. ### 🟢 KI-8 — conftest accessed non-existent field ✅ FIXED IN DAY 2 +## Status: 7/8 closed + +Remaining: KI-7 (day 6 — migrate the four agents to the new contract). +After day 6, day 7 un-xfails the smoke tests for full end-to-end +regression protection. + ## Progress tracker | Day | Issue(s) addressed | Status | @@ -45,6 +47,6 @@ Actual method is `call()`. **Fixed by:** Day 6. | 2 | KI-8 | ✅ | | 3 | KI-2, KI-5 | ✅ | | 4 | KI-3, KI-4, KI-6 | ✅ | -| 5 | KI-1 | ⏳ | -| 6 | KI-7, KI-4 (llm.chat)| ⏳ | -| 7 | All checked | ⏳ | +| 5 | KI-1 | ✅ | +| 6 | KI-7 + agent migration | ⏳ | +| 7 | un-xfail smoke tests | ⏳ | diff --git a/tests/unit/test_orchestrator.py b/tests/unit/test_orchestrator.py index ea90d68..69fcce1 100644 --- a/tests/unit/test_orchestrator.py +++ b/tests/unit/test_orchestrator.py @@ -56,11 +56,9 @@ def test_dry_run_kb_has_dry_run_keys(): def test_dry_run_authorized_scope(): - orch = Orchestrator( - authorized_scope=["10.0.0.0/24"], - dry_run=True, - ) - session = orch.run("10.0.0.1") + # authorized_scope moved from constructor to run() in day 5 + orch = Orchestrator(dry_run=True) + session = orch.run("10.0.0.1", authorized_scope=["10.0.0.0/24"]) assert "10.0.0.0/24" in session.authorized_scope diff --git a/tests/unit/test_orchestrator_config.py b/tests/unit/test_orchestrator_config.py new file mode 100644 index 0000000..fe26046 --- /dev/null +++ b/tests/unit/test_orchestrator_config.py @@ -0,0 +1,78 @@ +""" +Tests for Orchestrator + CyberAIConfig integration — day 5 of STANDOFF. + +Covers the KI-1 fix: orchestrator accepts config, builds llm lazily, +and the CLI wiring works in dry-run. +""" +from __future__ import annotations + +import pytest +from click.testing import CliRunner + +from cyberai.__main__ import cli +from cyberai.core.config import CyberAIConfig +from cyberai.core.orchestrator import Orchestrator +from cyberai.core.scan_session import ScanState + + +# ── Orchestrator + config ───────────────────────────────────────────── + + +def test_orchestrator_accepts_config(): + orch = Orchestrator(config=CyberAIConfig(), dry_run=True) + assert orch.config is not None + + +def test_orchestrator_defaults_config_when_omitted(): + """Orchestrator() with no config should still build a default one.""" + orch = Orchestrator(dry_run=True) + assert isinstance(orch.config, CyberAIConfig) + + +def test_orchestrator_llm_is_none_in_dry_run(): + """Dry-run must never construct an LLM client (no API key needed).""" + orch = Orchestrator(config=CyberAIConfig(), dry_run=True) + assert orch.llm is None + + +def test_orchestrator_run_returns_completed_session(): + orch = Orchestrator(config=CyberAIConfig(), dry_run=True) + session = orch.run("127.0.0.1") + assert session.state == ScanState.COMPLETED + assert len(session.phases) == 4 + + +def test_orchestrator_run_accepts_scope(): + orch = Orchestrator(config=CyberAIConfig(), dry_run=True) + session = orch.run("10.0.0.1", authorized_scope=["10.0.0.0/24"]) + assert "10.0.0.0/24" in session.authorized_scope + + +# ── CLI wiring ──────────────────────────────────────────────────────── + + +def test_cli_scan_dry_run_exits_zero(): + runner = CliRunner() + result = runner.invoke(cli, ["scan", "127.0.0.1", "--dry-run"]) + assert result.exit_code == 0, result.output + + +def test_cli_scan_dry_run_with_scope(): + runner = CliRunner() + result = runner.invoke( + cli, ["scan", "10.0.0.1", "--dry-run", "--scope", "10.0.0.0/24"] + ) + assert result.exit_code == 0, result.output + + +def test_cli_scan_reports_findings_count(): + runner = CliRunner() + result = runner.invoke(cli, ["scan", "127.0.0.1", "--dry-run"]) + assert "Findings:" in result.output + + +def test_cli_status_works(): + runner = CliRunner() + result = runner.invoke(cli, ["status"]) + assert result.exit_code == 0 + assert "Provider" in result.output