orithos-cli 0.1.0__py3-none-any.whl

orithos_cli/scan.py

@@ -0,0 +1,758 @@
+"""Scan commands for Orithos CLI."""
+
+from __future__ import annotations
+
+import json
+import math
+import time
+import httpx
+import click
+
+from orithos_cli.config import get_config, fmt_http_error
+from orithos_cli.output import (
+    FindingOutput,
+    ScanStatusOutput,
+    format_summary,
+    format_json,
+    format_sarif,
+    format_junit,
+)
+
+
+DEFAULT_PROBES = [
+    "prompt_injection",
+    "tool_misuse",
+    "scope_escalation",
+    "exfil_attempt",
+]
+
+ORITHOS_VERSION = "0.1.0"
+
+
+def _s(
+    text: str = "", fg: str | None = None, dim: bool = False, bold: bool = False
+) -> str:
+    return click.style(text, fg=fg, dim=dim, bold=bold)
+
+
+def _sev_color(sev: str) -> str:
+    if sev in ("critical", "high"):
+        return "red"
+    if sev == "medium":
+        return "yellow"
+    if sev == "low":
+        return "green"
+    return None
+
+
+def _fmt(s: str) -> str:
+    return s.replace("\n", "\\n")[:80]
+
+
+def _resolve_agent(cfg, agent: str) -> dict | None:
+    """Resolve agent name or ID. Tries ID first, then name search."""
+    try:
+        r = httpx.get(
+            f"{cfg.api_url}/v1/agents/{agent}",
+            headers=cfg.auth_headers(),
+            timeout=cfg.timeout,
+        )
+        if r.status_code == 200:
+            return r.json()
+    except httpx.RequestError:
+        pass
+    try:
+        r = httpx.get(
+            f"{cfg.api_url}/v1/agents",
+            headers=cfg.auth_headers(),
+            timeout=cfg.timeout,
+        )
+        if r.status_code == 200:
+            agents = r.json()
+            if isinstance(agents, list):
+                for a in agents:
+                    if a.get("name") == agent or a.get("id") == agent:
+                        return a
+    except httpx.RequestError:
+        pass
+    return None
+
+
+def _parse_tools(tools_str: str) -> list[dict]:
+    try:
+        parsed = json.loads(tools_str)
+        if isinstance(parsed, list):
+            return parsed
+    except (json.JSONDecodeError, TypeError):
+        pass
+    return []
+
+
+def _tool_summary(tools: list[dict]) -> str:
+    parts = []
+    for t in tools[:4]:
+        name = t.get("name", "?")
+        params = t.get("parameters", [])
+        access = "read"
+        if any(
+            p.get("name", "").lower() in ("amount", "value", "price") for p in params
+        ):
+            access = "write·financial"
+        elif any(
+            p.get("name", "").lower() in ("password", "key", "secret") for p in params
+        ):
+            access = "write·sensitive"
+        elif params:
+            access = "write"
+        parts.append(f"{name} ({access})")
+    if len(tools) > 4:
+        parts.append(f"+{len(tools) - 4} more")
+    return " · ".join(parts) if parts else "none declared"
+
+
+def _print_prescan(agent: dict) -> None:
+    prompt = agent.get("system_prompt") or agent.get("prompt_preview") or ""
+    tokens = max(1, len(prompt) // 4)
+    constraints = agent.get("constraints_count", 0) or max(
+        1, prompt.count("must") + prompt.count("never") + prompt.count("always")
+    )
+    identifiers = agent.get("identifiers_count", 0) or sum(
+        1 for w in prompt.split() if w.startswith("CSB-") or w.startswith("ID:")
+    )
+
+    tools = _parse_tools(agent.get("tools_config", "[]"))
+    memory = agent.get("memory_sources", "[]")
+    guardrails = agent.get("guardrails", "[]")
+
+    try:
+        mem_list = json.loads(memory) if isinstance(memory, str) else memory
+    except (json.JSONDecodeError, TypeError):
+        mem_list = []
+    try:
+        gd_list = json.loads(guardrails) if isinstance(guardrails, str) else guardrails
+    except (json.JSONDecodeError, TypeError):
+        gd_list = []
+
+    click.echo(_s(f"\u2726 Orithos v{ORITHOS_VERSION}", fg="green"))
+    click.echo(_s("  Loading agent configuration...", dim=True))
+    click.echo(
+        _s(
+            f"  \u21b3 System prompt: {tokens} tokens \u00b7 {constraints} constraints extracted \u00b7 {identifiers} identifiers flagged",
+            dim=True,
+        )
+    )
+
+    tool_line = _tool_summary(tools)
+    if tool_line:
+        click.echo(_s(f"  \u21b3 Tools: {tool_line}", dim=True))
+
+    if mem_list:
+        for m in mem_list[:2]:
+            provider = m.get("provider", "?")
+            doc_count = m.get("doc_count", "?")
+            risk = m.get("risk", "")
+            click.echo(
+                _s(
+                    f"  \u21b3 RAG: {m.get('name', 'knowledge-base')} \u00b7 {provider} \u00b7 {doc_count} docs"
+                    + (f" ({risk} risk flagged)" if risk else ""),
+                    dim=True,
+                )
+            )
+
+    for g in gd_list[:2]:
+        gtype = g.get("type", "?")
+        stage = g.get("stage", "both")
+        click.echo(_s(f"  \u21b3 Guardrails: {gtype} ({stage})", dim=True))
+
+    total_probes = agent.get("probe_count", 0) or 245
+    click.echo(_s(f"  Generating {total_probes} agent-targeted probes...", dim=True))
+    click.echo(_s(f"  \u21b3 {total_probes} probes \u00b7 mixed categories", dim=True))
+    click.echo(_s("\u2726 Scan executing \u00b7 session-isolated", fg="green"))
+
+
+def _poll_with_progress(
+    status_url: str,
+    cfg,
+    poll_interval: int,
+    max_wait: int,
+    total_probes: int,
+) -> dict:
+    start_time = time.monotonic()
+    last_probe = ""
+    while True:
+        elapsed = time.monotonic() - start_time
+        if elapsed > max_wait:
+            click.echo(
+                _s(f"Timed out after {max_wait}s.", fg="red", bold=True),
+                err=True,
+            )
+            raise SystemExit(1)
+        try:
+            r = httpx.get(status_url, headers=cfg.auth_headers(), timeout=cfg.timeout)
+            r.raise_for_status()
+        except httpx.RequestError as exc:
+            click.echo(_s(f"Warning: status poll failed: {exc}", fg="yellow"), err=True)
+            time.sleep(poll_interval)
+            continue
+        scan = r.json()
+        status = scan["status"]
+        done = scan.get("completed_tests", 0)
+        total = scan.get("total_tests", total_probes) or total_probes
+        pct = math.floor(done / total * 100) if total > 0 else 0
+
+        probe_name = scan.get("current_probe_name", "")
+        stage = scan.get("current_probe_stage", "")
+
+        if status in ("completed", "failed"):
+            mins = int(elapsed // 60)
+            secs = int(elapsed % 60)
+            if status == "completed":
+                click.echo(
+                    _s(
+                        f"\n\u2726 Scan complete \u00b7 {done}/{total} probes \u00b7 {mins}m {secs}s",
+                        fg="green",
+                        bold=True,
+                    )
+                )
+            else:
+                click.echo(
+                    _s(
+                        f"\n\u2716 Scan failed \u00b7 {done}/{total} probes \u00b7 {mins}m {secs}s",
+                        fg="red",
+                        bold=True,
+                    )
+                )
+            return scan
+
+        if probe_name and probe_name != last_probe:
+            last_probe = probe_name
+            click.echo(
+                _s(
+                    f"  \u26a0  {_fmt(probe_name)} \u2014 {stage or 'running'}",
+                    fg="yellow",
+                )
+            )
+
+        bar_len = 20
+        filled = round(bar_len * pct / 100)
+        bar = _s("\u2588" * filled, fg="green") + _s(
+            "\u2591" * (bar_len - filled), dim=True
+        )
+        pct_str = (
+            _s(f"{pct}%", dim=True) if pct < 100 else _s("100%", fg="green", bold=True)
+        )
+        click.echo(
+            f"\r  [{bar}] {pct_str} \u00b7 {_s(f'{done}/{total}', dim=True)} \u00b7 [{_s(f'{elapsed:.0f}s', dim=True)}]",
+            nl=False,
+        )
+        time.sleep(poll_interval)
+
+
+def _print_completion(scan: dict, findings: list[FindingOutput]) -> None:
+    risk = scan.get("severity", "unknown").upper()
+    total = len(findings)
+    by_sev: dict[str, int] = {}
+    for f in findings:
+        by_sev[f.severity] = by_sev.get(f.severity, 0) + 1
+
+    crit = by_sev.get("critical", 0)
+    high = by_sev.get("high", 0)
+    med = by_sev.get("medium", 0)
+    low = by_sev.get("low", 0)
+
+    parts = []
+    if crit:
+        parts.append(f"{crit} critical")
+    if high:
+        parts.append(f"{high} high")
+    if med:
+        parts.append(f"{med} medium")
+    if low:
+        parts.append(f"{low} low")
+    findings_str = " \u00b7  ".join(parts) if parts else "0 findings"
+
+    risk_c = _sev_color(risk.lower())
+    click.echo(
+        _s("  Risk score: ", dim=True)
+        + _s(risk, fg=risk_c, bold=True)
+        + _s(f"  \u00b7  {findings_str}", dim=True)
+    )
+    click.echo(_s("  Scan ID:    ", dim=True) + _s(scan["id"], dim=True))
+
+    if findings:
+        click.echo("")
+        for f in findings[:3]:
+            fc = _sev_color(f.severity)
+            click.echo(
+                f"  [{_s(f.severity.upper(), fg=fc, bold=True)}] {_fmt(f.attack_path)}"
+            )
+            if f.recommendation:
+                click.echo(_s(f"       {_fmt(f.recommendation)}", dim=True))
+        if total > 3:
+            click.echo(_s(f"  ... and {total - 3} more findings", dim=True))
+    else:
+        click.echo(_s("  No findings \u2014 agent is secure.", fg="green"))
+
+
+def _run_scan(
+    agent: str,
+    probes: str,
+    output_format: str,
+    poll_interval: int,
+    max_wait: int,
+    output_file: str | None,
+    connection_id: str | None = None,
+    config: dict | None = None,
+) -> None:
+    cfg = get_config()
+
+    agent_obj = _resolve_agent(cfg, agent)
+    if not agent_obj:
+        click.echo(f"Error: agent '{agent}' not found", err=True)
+        raise SystemExit(1)
+
+    agent_id = agent_obj["id"]
+    total_probes = agent_obj.get("probe_count", 0) or 245
+
+    _print_prescan(agent_obj)
+
+    # Auto-use preferred connection if no explicit connection given
+    if not connection_id and agent_obj.get("preferred_connection_id"):
+        connection_id = agent_obj["preferred_connection_id"]
+        click.echo(
+            click.style(f"  Using preferred connection: {connection_id[:8]}…", dim=True)
+        )
+
+    probe_list = [p.strip() for p in probes.split(",") if p.strip()]
+    body: dict = {"agent_id": agent_id, "probe_set": probe_list}
+    if connection_id:
+        body["connection_id"] = connection_id
+    if config:
+        body["config"] = config
+    try:
+        r = httpx.post(
+            f"{cfg.api_url}/v1/scans",
+            json=body,
+            headers=cfg.auth_headers(),
+            timeout=cfg.timeout,
+        )
+        r.raise_for_status()
+    except httpx.HTTPStatusError as exc:
+        click.echo(f"Error creating scan: {fmt_http_error(exc)}", err=True)
+        raise SystemExit(1)
+    except httpx.RequestError as exc:
+        click.echo(f"Connection error: {exc}", err=True)
+        raise SystemExit(1)
+
+    scan_data = r.json()
+    scan_id = scan_data["scan_id"]
+
+    # Poll with progress
+    status_url = f"{cfg.api_url}/v1/scans/{scan_id}"
+    scan_status = _poll_with_progress(
+        status_url, cfg, poll_interval, max_wait, total_probes
+    )
+
+    if scan_status["status"] == "failed":
+        click.echo("", err=True)
+        raise SystemExit(1)
+
+    # Fetch findings
+    try:
+        fr = httpx.get(
+            f"{status_url}/findings",
+            headers=cfg.auth_headers(),
+            timeout=cfg.timeout,
+        )
+        fr.raise_for_status()
+    except httpx.RequestError as exc:
+        click.echo(f"Error fetching findings: {exc}", err=True)
+        raise SystemExit(1)
+
+    findings_data = fr.json()
+    # GET /v1/scans/{id}/findings returns a PaginatedFindingsResponse
+    # envelope ({"items": [...], "total": ...}), not a bare list.
+    if isinstance(findings_data, dict):
+        findings_data = findings_data.get("items", [])
+    findings = [FindingOutput(**f) for f in findings_data]
+
+    if output_format == "summary":
+        _print_completion(scan_status, findings)
+        return
+
+    # Structured output
+    scan_out = ScanStatusOutput(
+        id=scan_id,
+        status=scan_status["status"],
+        severity=scan_status.get("severity"),
+        total_tests=scan_status.get("total_tests", 0),
+        completed_tests=scan_status.get("completed_tests", 0),
+    )
+
+    if output_format == "json":
+        output = format_json(findings, scan_out)
+    elif output_format == "sarif":
+        output = format_sarif(findings, scan_out, cfg.api_url)
+    elif output_format == "junit":
+        output = format_junit(findings, scan_out)
+    else:
+        output = format_summary(findings, scan_out)
+
+    if output_file:
+        try:
+            with open(output_file, "w") as f:
+                f.write(output)
+            click.echo(f"Output written to {output_file}")
+        except OSError as exc:
+            click.echo(f"Failed to write output file: {exc}", err=True)
+            raise SystemExit(1)
+    else:
+        click.echo(output)
+
+
+@click.group(name="scan", invoke_without_command=True)
+@click.pass_context
+@click.option("--agent", "-a", default="", help="Agent name or ID to scan")
+@click.option(
+    "--connection",
+    "-c",
+    default="",
+    help="Connection name or ID (default: agent's preferred connection)",
+)
+@click.option(
+    "--probes",
+    default=",".join(DEFAULT_PROBES),
+    help="Comma-separated probe names (default: all probes)",
+    show_default=False,
+)
+@click.option(
+    "--format",
+    "output_format",
+    type=click.Choice(["summary", "json", "sarif", "junit"], case_sensitive=False),
+    default="summary",
+    help="Output format",
+)
+@click.option(
+    "--poll-interval",
+    default=2,
+    type=int,
+    help="Seconds between status polls (default: 2)",
+    hidden=True,
+)
+@click.option(
+    "--max-wait",
+    default=600,
+    type=int,
+    help="Max seconds to wait for scan completion (default: 600)",
+)
+@click.option(
+    "--output-file",
+    "output_file",
+    type=click.Path(dir_okay=False, writable=True),
+    default=None,
+    help="Write output to file instead of stdout",
+)
+def scan_group(
+    ctx: click.Context,
+    agent: str,
+    connection: str,
+    probes: str,
+    output_format: str,
+    poll_interval: int,
+    max_wait: int,
+    output_file: str | None,
+) -> None:
+    """Scan agents for security vulnerabilities.
+
+    Runs a full security scan against an agent and streams live results.
+
+    Examples:
+
+        orithos scan --agent "Customer Support Agent"
+
+        orithos scan --agent "Customer Support Agent" --connection staging
+
+        orithos scan --agent 12c2733b --format json
+    """
+    if ctx.invoked_subcommand is not None:
+        return
+
+    if not agent:
+        click.echo(ctx.get_help())
+        return
+
+    connection_id = connection or None
+    _run_scan(
+        agent,
+        probes,
+        output_format,
+        poll_interval,
+        max_wait,
+        output_file,
+        connection_id,
+    )
+
+
+@scan_group.command("new")
+def scan_new() -> None:
+    """Interactive scan wizard — configure and launch."""
+    from orithos_cli.wizard import interactive_wizard
+
+    interactive_wizard()
+
+
+@scan_group.command("run")
+def scan_run() -> None:
+    """Interactive scan wizard — alias for 'scan new'."""
+    from orithos_cli.wizard import interactive_wizard
+
+    interactive_wizard()
+
+
+@scan_group.command("status")
+@click.argument("scan_id")
+def scan_status(scan_id: str) -> None:
+    """Check the status of a scan."""
+    cfg = get_config()
+    try:
+        r = httpx.get(
+            f"{cfg.api_url}/v1/scans/{scan_id}",
+            headers=cfg.auth_headers(),
+            timeout=cfg.timeout,
+        )
+        r.raise_for_status()
+    except httpx.HTTPStatusError as exc:
+        click.echo(f"Scan not found: {exc.response.status_code}", err=True)
+        raise SystemExit(1)
+    except httpx.RequestError as exc:
+        click.echo(f"Connection error: {exc}", err=True)
+        raise SystemExit(1)
+
+    data = r.json()
+    click.echo(f"Scan:      {data['id']}")
+    click.echo(f"Status:    {data['status']}")
+    click.echo(f"Severity:  {data.get('severity') or 'N/A'}")
+    click.echo(
+        f"Tests:     {data.get('completed_tests', 0)}/{data.get('total_tests', 0)}"
+    )
+    if data.get("started_at"):
+        click.echo(f"Started:   {data['started_at']}")
+    if data.get("completed_at"):
+        click.echo(f"Completed: {data['completed_at']}")
+
+
+@scan_group.command("findings")
+@click.argument("scan_id")
+def scan_findings(scan_id: str) -> None:
+    """Print findings for a completed scan."""
+    cfg = get_config()
+    try:
+        r = httpx.get(
+            f"{cfg.api_url}/v1/scans/{scan_id}/findings",
+            headers=cfg.auth_headers(),
+            timeout=cfg.timeout,
+        )
+        r.raise_for_status()
+    except httpx.HTTPStatusError as exc:
+        click.echo(f"Error: {exc.response.status_code}", err=True)
+        raise SystemExit(1)
+    except httpx.RequestError as exc:
+        click.echo(f"Connection error: {exc}", err=True)
+        raise SystemExit(1)
+
+    findings = r.json()
+    if not findings:
+        click.echo("No findings.")
+        return
+    for f in findings:
+        click.echo(
+            f"[{f['severity'].upper()}] {f.get('attack_path', f.get('name', '?'))}"
+        )
+        if f.get("recommendation"):
+            click.echo(f"  {f['recommendation']}")
+        click.echo("")
+
+
+@scan_group.command("rerun")
+@click.argument("scan_id")
+@click.option(
+    "--scope",
+    "rerun_scope",
+    type=click.Choice(["same", "failed_only", "passed_only"], case_sensitive=False),
+    default="same",
+    help="Probe scope for rerun",
+)
+@click.option(
+    "--format",
+    "output_format",
+    type=click.Choice(["summary", "json", "sarif", "junit"], case_sensitive=False),
+    default="summary",
+    help="Output format",
+)
+@click.option(
+    "--output-file",
+    type=click.Path(dir_okay=False, writable=True),
+    default=None,
+    help="Write output to file instead of stdout",
+)
+@click.option(
+    "--poll-interval",
+    default=2,
+    type=int,
+    help="Seconds between status polls (default: 2)",
+    hidden=True,
+)
+@click.option(
+    "--max-wait",
+    default=600,
+    type=int,
+    help="Max seconds to wait for scan completion (default: 600)",
+)
+def scan_rerun(
+    scan_id: str,
+    rerun_scope: str,
+    output_format: str,
+    output_file: str | None,
+    poll_interval: int,
+    max_wait: int,
+) -> None:
+    """Rerun an existing scan with scope options."""
+    cfg = get_config()
+
+    try:
+        r = httpx.post(
+            f"{cfg.api_url}/v1/scans/{scan_id}/rerun",
+            params={"scope": rerun_scope},
+            headers=cfg.auth_headers(),
+            timeout=cfg.timeout,
+        )
+        r.raise_for_status()
+    except httpx.HTTPStatusError as exc:
+        click.echo(f"Error rerunning scan: {fmt_http_error(exc)}", err=True)
+        raise SystemExit(1)
+    except httpx.RequestError as exc:
+        click.echo(f"Connection error: {exc}", err=True)
+        raise SystemExit(1)
+
+    rerun_data = r.json()
+    new_scan_id = rerun_data["scan_id"]
+    click.echo(f"Rerun created: {new_scan_id} (originated from {scan_id})")
+
+    status_url = f"{cfg.api_url}/v1/scans/{new_scan_id}"
+    scan_status = _poll_with_progress(status_url, cfg, poll_interval, max_wait, 0)
+
+    if scan_status["status"] == "failed":
+        raise SystemExit(1)
+
+    try:
+        fr = httpx.get(
+            f"{status_url}/findings",
+            headers=cfg.auth_headers(),
+            timeout=cfg.timeout,
+        )
+        fr.raise_for_status()
+    except httpx.RequestError as exc:
+        click.echo(f"Error fetching findings: {exc}", err=True)
+        raise SystemExit(1)
+
+    findings_data = fr.json()
+    # GET /v1/scans/{id}/findings returns a PaginatedFindingsResponse
+    # envelope ({"items": [...], "total": ...}), not a bare list.
+    if isinstance(findings_data, dict):
+        findings_data = findings_data.get("items", [])
+    findings = [FindingOutput(**f) for f in findings_data]
+    scan_out = ScanStatusOutput(
+        id=new_scan_id,
+        status=scan_status["status"],
+        severity=scan_status.get("severity"),
+        total_tests=scan_status.get("total_tests", 0),
+        completed_tests=scan_status.get("completed_tests", 0),
+    )
+
+    if output_format == "json":
+        output = format_json(findings, scan_out)
+    elif output_format == "sarif":
+        output = format_sarif(findings, scan_out, cfg.api_url)
+    elif output_format == "junit":
+        output = format_junit(findings, scan_out)
+    else:
+        output = format_summary(findings, scan_out)
+
+    if output_file:
+        try:
+            with open(output_file, "w") as f:
+                f.write(output)
+            click.echo(f"Output written to {output_file}")
+        except OSError as exc:
+            click.echo(f"Failed to write output file: {exc}", err=True)
+            raise SystemExit(1)
+    else:
+        click.echo(output)
+
+
+@scan_group.command("report")
+@click.argument("scan_id")
+@click.option(
+    "--format",
+    "output_format",
+    type=click.Choice(["summary", "json"], case_sensitive=False),
+    default="summary",
+    help="Output format",
+)
+def scan_report(scan_id: str, output_format: str) -> None:
+    """Generate a full security report for a completed scan."""
+    cfg = get_config()
+    try:
+        r = httpx.get(
+            f"{cfg.api_url}/v1/scans/{scan_id}/report",
+            headers=cfg.auth_headers(),
+            timeout=cfg.timeout,
+        )
+        r.raise_for_status()
+    except httpx.HTTPStatusError as exc:
+        click.echo(f"Scan not found: {exc.response.status_code}", err=True)
+        raise SystemExit(1)
+    except httpx.RequestError as exc:
+        click.echo(f"Connection error: {exc}", err=True)
+        raise SystemExit(1)
+
+    report = r.json()
+    if output_format == "json":
+        click.echo(json.dumps(report, indent=2))
+        return
+
+    click.echo("Orithos Security Report")
+    click.echo(f"{'=' * 50}")
+    click.echo(f"Scan ID:    {report['scan_id']}")
+    click.echo(f"Status:     {report['status']}")
+    click.echo(f"Severity:   {report.get('severity') or 'N/A'}")
+    click.echo(f"Agent ID:   {report['agent_id']}")
+
+    summary = report.get("summary", {})
+    click.echo("")
+    click.echo("Summary:")
+    click.echo(f"  Total tests:   {summary.get('total', 0)}")
+    click.echo(f"  Completed:     {summary.get('completed', 0)}")
+    click.echo(f"  High:          {summary.get('high', 0)}")
+    click.echo(f"  Medium:        {summary.get('medium', 0)}")
+    click.echo(f"  Low:           {summary.get('low', 0)}")
+
+    findings = report.get("findings", [])
+    click.echo("")
+    click.echo(f"Findings ({len(findings)}):")
+    if findings:
+        for f in findings:
+            click.echo(
+                f"  [{f['severity'].upper()}] {f.get('attack_path', f.get('name', '?'))}"
+            )
+            if f.get("recommendation"):
+                click.echo(f"    {f['recommendation']}")
+    else:
+        click.echo("  No findings \u2014 agent is secure.")
+
+    recommendations = report.get("recommendations", [])
+    if recommendations:
+        click.echo("")
+        click.echo("Top Recommendations:")
+        for i, rec in enumerate(recommendations[:5], 1):
+            click.echo(f"  {i}. {rec}")