pentest.py

#!/usr/bin/env python3
"""pentestkit CLI — drive a multi-agent, scope-guarded penetration test.

Usage:
  python pentest.py run    --engagement engagements/example.yaml [--yes] [--resume]
  python pentest.py check  --engagement engagements/example.yaml --url https://t/x
  python pentest.py scopes
"""
from __future__ import annotations

import argparse
import asyncio
import sys
import time
from pathlib import Path

from pentestkit.config import load_engagement, load_env, require_api_key
from pentestkit.guardrails import ScopeGuard
from pentestkit.scopes import ALL_SCOPES

ROOT = Path(__file__).resolve().parent


# ---- pretty console output -------------------------------------------------
class C:
    G = "\033[32m"; Y = "\033[33m"; R = "\033[31m"; B = "\033[34m"
    DIM = "\033[2m"; BOLD = "\033[1m"; X = "\033[0m"


def _emit(phase: str, channel: str, message: str) -> None:
    ts = time.strftime("%H:%M:%S")
    color = {"phase": C.BOLD + C.B, "tool": C.DIM, "info": C.G,
             "text": C.DIM, "error": C.R}.get(channel, "")
    tag = f"{color}{phase:>14}{C.X}"
    if channel == "tool":
        print(f"{C.DIM}{ts}{C.X} {tag} {C.DIM}· {message}{C.X}")
    elif channel == "text":
        snippet = message[:160].replace("\n", " ")
        print(f"{C.DIM}{ts} {tag}   {snippet}{C.X}")
    else:
        print(f"{C.DIM}{ts}{C.X} {tag} {color}{message}{C.X}")


def _confirm_authorization(eng, assume_yes: bool) -> bool:
    print(f"\n{C.BOLD}AUTHORIZATION CHECK{C.X}")
    print(f"  Engagement : {eng.name}")
    print(f"  Target     : {eng.base_url}")
    print(f"  Authorized : {eng.authorized_by or C.R + 'NOT SPECIFIED' + C.X}")
    print(f"  Reference  : {eng.authorization_ref or C.R + 'NOT SPECIFIED' + C.X}")
    print(f"  In-scope   : hosts={eng.scope.hosts or '—'} cidrs={eng.scope.cidrs or '—'} "
          f"ports={eng.scope.ports or 'any'}")
    print(f"  Scopes     : {', '.join(eng.scopes_to_run)}")
    print(f"  RoE        : {len(eng.rules_of_engagement)} rule(s)")
    if assume_yes:
        print(f"{C.Y}  --yes supplied; proceeding.{C.X}")
        return True
    ans = input(f"\n{C.Y}Confirm you are AUTHORIZED to test the above. Type 'I CONFIRM': {C.X}")
    return ans.strip() == "I CONFIRM"


# ---- commands --------------------------------------------------------------
async def cmd_run(args) -> int:
    load_env(args.env)
    try:
        require_api_key()
    except RuntimeError as e:
        print(f"{C.R}{e}{C.X}")
        return 2
    eng = load_engagement(args.engagement)
    if not eng.base_url and not eng.scope.hosts and not eng.scope.cidrs:
        print(f"{C.R}Engagement has no target/scope defined.{C.X}")
        return 2
    if not _confirm_authorization(eng, args.yes):
        print(f"{C.R}Authorization not confirmed. Aborting.{C.X}")
        return 1

    out_dir = Path(args.out) if args.out else (ROOT / "output" /
              f"{_slug(eng.name)}-{time.strftime('%Y%m%d-%H%M%S')}")
    out_dir.mkdir(parents=True, exist_ok=True)
    # import here so a missing SDK gives a clean message only when actually running
    from pentestkit.pipeline.orchestrator import Orchestrator

    orch = Orchestrator(eng, out_dir, on_event=_emit, resume=args.resume,
                        run_kind=getattr(args, "run_kind", "pentest") or "pentest",
                        run_group=getattr(args, "run_group", "") or "",
                        benchmark_id=getattr(args, "benchmark_id", "") or "")
    orch.commit_worklog = bool(getattr(args, "commit_worklog", False))

    httpd = None
    if args.ui:
        from pentestkit.web import serve
        import threading
        httpd = serve(out_dir, port=args.ui_port)
        threading.Thread(target=httpd.serve_forever, daemon=True).start()
        print(f"\n{C.BOLD}{C.B}Live UI: http://127.0.0.1:{args.ui_port}{C.X}  "
              f"{C.DIM}(open in a browser){C.X}")

    print(f"\n{C.BOLD}Starting engagement.{C.X} Artifacts -> {out_dir}\n")
    t0 = time.time()
    try:
        report_path = await orch.run(skip_plan=args.skip_plan)
    finally:
        await orch.close()
    dt = time.time() - t0
    print(f"\n{C.G}{C.BOLD}Done in {dt:.0f}s.{C.X}")
    print(f"  Report      : {report_path}")
    print(f"  Context/KB  : {out_dir / 'context.json'}")
    print(f"  Run summary : {out_dir / 'run_summary.md'}")
    print(f"  Evidence    : {out_dir / 'evidence'}/")
    print(f"  Total cost  : ${orch.total_cost():.4f}")
    if httpd is not None:
        print(f"\n{C.B}UI still live at http://127.0.0.1:{args.ui_port} — "
              f"Ctrl-C to exit.{C.X}")
        try:
            while True:
                time.sleep(3600)
        except KeyboardInterrupt:
            httpd.shutdown()
    return 0


async def cmd_fleet(args) -> int:
    load_env(args.env)
    try:
        require_api_key()
    except RuntimeError as e:
        print(f"{C.R}{e}{C.X}")
        return 2
    from pentestkit.config import load_fleet
    from pentestkit.fleet import plan_fleet, run_fleet
    from pentestkit.web import serve_fleet

    fleet = load_fleet(args.fleet)
    if args.max_parallel:
        fleet.max_parallel_targets = args.max_parallel
    if args.dashboard_port:
        fleet.dashboard_port = args.dashboard_port

    fleet_dir = Path(args.out) if args.out else (ROOT / "output" /
                f"fleet-{time.strftime('%Y%m%d-%H%M%S')}")
    fleet_dir.mkdir(parents=True, exist_ok=True)
    plans = plan_fleet(fleet, fleet_dir)

    print(f"\n{C.BOLD}FLEET: {fleet.name}{C.X}")
    print(f"  Targets ({len(plans)}), {fleet.max_parallel_targets} in parallel:")
    for p in plans:
        print(f"    • {p['name']:<32} -> {p['target']}")
    print(f"  Output: {fleet_dir}")
    if not args.yes:
        ans = input(f"\n{C.Y}Confirm you are AUTHORIZED to test ALL targets above. "
                    f"Type 'I CONFIRM': {C.X}")
        if ans.strip() != "I CONFIRM":
            print(f"{C.R}Authorization not confirmed. Aborting.{C.X}")
            return 1

    httpd = serve_fleet(fleet_dir, port=fleet.dashboard_port,
                        interactive=True, env_path=args.env)
    import threading
    threading.Thread(target=httpd.serve_forever, daemon=True).start()
    print(f"\n{C.BOLD}{C.B}Fleet dashboard: http://127.0.0.1:{fleet.dashboard_port}{C.X}  "
          f"{C.DIM}(+ Add target to launch more from the UI){C.X}\n")

    t0 = time.time()
    results = await run_fleet(fleet, plans, on_log=lambda m: print(f"  {C.DIM}{m}{C.X}"))
    dt = time.time() - t0
    ok = sum(1 for r in results if r["rc"] == 0)
    print(f"\n{C.G}{C.BOLD}Fleet complete in {dt:.0f}s — {ok}/{len(results)} succeeded.{C.X}")
    print(f"  Artifacts: {fleet_dir}/<target>/")
    print(f"\n{C.B}Dashboard still live at http://127.0.0.1:{fleet.dashboard_port} — "
          f"Ctrl-C to exit.{C.X}")
    try:
        while True:
            time.sleep(3600)
    except KeyboardInterrupt:
        httpd.shutdown()
    return 0


def cmd_console(args) -> int:
    load_env(args.env)
    try:
        require_api_key()
    except RuntimeError as e:
        print(f"{C.R}{e}{C.X}")
        return 2
    fleet_dir = (Path(args.fleet_dir) if args.fleet_dir else
                 ROOT / "output" / f"console-{time.strftime('%Y%m%d-%H%M%S')}")
    fleet_dir.mkdir(parents=True, exist_ok=True)
    results_root = Path(args.results_dir) if args.results_dir else \
        (ROOT / "benchmarks" / "xbow" / "results")
    from pentestkit.web import serve_console
    httpd = serve_console(fleet_dir, results_root, port=args.port,
                          interactive=True, env_path=args.env)
    print(f"{C.BOLD}{C.B}pentestkit console -> http://127.0.0.1:{args.port}{C.X}")
    print(f"{C.DIM}Tabs: Fleet (add/launch targets) · Benchmarks (latest run) · "
          f"Traces (run/agent DB).{C.X}")
    print(f"{C.DIM}Fleet dir: {fleet_dir}   Benchmarks: {results_root}{C.X}")
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        httpd.shutdown()
    return 0


def cmd_fleet_console(args) -> int:
    load_env(args.env)
    try:
        require_api_key()
    except RuntimeError as e:
        print(f"{C.R}{e}{C.X}")
        return 2
    fleet_dir = (Path(args.dir) if args.dir else
                 ROOT / "output" / f"fleet-console-{time.strftime('%Y%m%d-%H%M%S')}")
    fleet_dir.mkdir(parents=True, exist_ok=True)
    from pentestkit.web import serve_fleet
    httpd = serve_fleet(fleet_dir, port=args.port, interactive=True, env_path=args.env)
    print(f"{C.BOLD}{C.B}Fleet console -> http://127.0.0.1:{args.port}{C.X}")
    print(f"{C.DIM}Serving {fleet_dir}. Use '+ Add target' to launch pentests. "
          f"Ctrl-C to stop.{C.X}")
    print(f"{C.Y}Only add targets you are authorized to test.{C.X}")
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        httpd.shutdown()
    return 0


async def cmd_ci(args) -> int:
    """Non-interactive run for CI: emits SARIF and exits non-zero on severe findings."""
    load_env(args.env)
    try:
        require_api_key()
    except RuntimeError as e:
        print(f"{C.R}{e}{C.X}")
        return 2
    eng = load_engagement(args.engagement)
    out_dir = Path(args.out) if args.out else (ROOT / "output" /
              f"ci-{_slug(eng.name)}-{time.strftime('%Y%m%d-%H%M%S')}")
    out_dir.mkdir(parents=True, exist_ok=True)
    from pentestkit.pipeline.orchestrator import Orchestrator
    from pentestkit.sarif import worst_severity, SEV_ORDER
    from pentestkit.models import FindingStatus
    orch = Orchestrator(eng, out_dir, on_event=_emit, run_kind="ci")
    try:
        await orch.run(skip_plan=args.skip_plan)
    finally:
        await orch.close()
    real = orch.ctx.findings_by_status(FindingStatus.VERIFIED, FindingStatus.SCORED,
                                       FindingStatus.REPORTED)
    sarif = out_dir / "report.sarif"
    print(f"\n{C.BOLD}CI run complete.{C.X}")
    print(f"  SARIF   : {sarif}")
    print(f"  Findings: {len(real)} confirmed")
    worst = worst_severity(real)
    threshold = SEV_ORDER.get((args.fail_on or "high").lower(), 3)
    rank_name = {v: k for k, v in SEV_ORDER.items()}
    if worst >= threshold and real:
        print(f"{C.R}FAIL: highest severity '{rank_name.get(worst)}' >= gate "
              f"'{args.fail_on}'.{C.X}")
        return 1
    print(f"{C.G}PASS: no findings at or above gate '{args.fail_on}'.{C.X}")
    return 0


async def cmd_finish(args) -> int:
    """Resume an interrupted run: finish verify/score/report on existing context."""
    load_env(args.env)
    try:
        require_api_key()
    except RuntimeError as e:
        print(f"{C.R}{e}{C.X}")
        return 2
    run_dir = Path(args.run_dir)
    if not (run_dir / "context.json").exists():
        print(f"{C.R}No context.json in {run_dir} — nothing to finish.{C.X}")
        return 2
    eng = load_engagement(args.engagement)
    from pentestkit.pipeline.orchestrator import Orchestrator
    orch = Orchestrator(eng, run_dir, on_event=_emit, resume=True)
    print(f"\n{C.BOLD}Finishing run.{C.X} {run_dir}\n")
    try:
        if not args.skip_verify:
            await orch.verify()       # no-op if no candidates remain
            await orch.stability()
        await orch.score()
        report_path = await orch.report()
        orch._write_summary()
        orch.live.finish()
    finally:
        await orch.close()
    print(f"\n{C.G}{C.BOLD}Done.{C.X}  Report: {report_path}")
    return 0


def cmd_serve(args) -> int:
    run_dir = Path(args.run_dir)
    if not run_dir.exists():
        print(f"{C.R}No such run dir: {run_dir}{C.X}")
        return 2
    if not (run_dir / "context.json").exists():
        print(f"{C.Y}Warning: {run_dir}/context.json not found yet — the UI will populate "
              f"once the run writes it.{C.X}")
    from pentestkit.web import serve
    httpd = serve(run_dir, port=args.port)
    print(f"{C.BOLD}{C.B}pentestkit UI -> http://127.0.0.1:{args.port}{C.X}")
    print(f"{C.DIM}Serving {run_dir}. Ctrl-C to stop.{C.X}")
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        httpd.shutdown()
    return 0


def cmd_check(args) -> int:
    eng = load_engagement(args.engagement)
    guard = ScopeGuard(eng.scope)
    d = guard.check_url(args.url)
    mark = f"{C.G}ALLOWED{C.X}" if d.allowed else f"{C.R}BLOCKED{C.X}"
    print(f"{mark}: {args.url}\n  reason: {d.reason}")
    return 0 if d.allowed else 1


def cmd_serve_bench(args) -> int:
    run_dir = Path(args.results_dir)
    if not run_dir.exists():
        print(f"{C.R}No such results dir: {run_dir}{C.X}")
        return 2
    from pentestkit.web import serve_benchmark
    httpd = serve_benchmark(run_dir, port=args.port)
    print(f"{C.BOLD}{C.B}pentestkit benchmark dashboard -> http://127.0.0.1:{args.port}{C.X}")
    print(f"{C.DIM}Serving {run_dir}. Ctrl-C to stop.{C.X}")
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        httpd.shutdown()
    return 0


def cmd_history(args) -> int:
    from pentestkit.db import Database, import_from_disk
    db = Database()
    if getattr(args, "import_disk", False):
        roots = [ROOT / "output", ROOT / "benchmarks" / "xbow" / "results"]
        n = import_from_disk(db, roots, log=lambda m: print(f"  {C.DIM}{m}{C.X}"))
        print(f"{C.G}Imported {n} run(s) from disk.{C.X}\n")
    if args.agent_id:
        msgs = db.list_messages(args.agent_id)
        if not msgs:
            print(f"{C.Y}No messages for agent {args.agent_id}.{C.X}")
            return 1
        for m in msgs:
            tag = f"{C.B}{m['type']:>14}{C.X}"
            if m["type"] == "tool_use":
                print(f"{tag}  {C.BOLD}{m['tool_name']}{C.X} {(m['tool_input'] or '')[:160]}")
            elif m["type"] == "tool_result":
                err = f"{C.R}[err]{C.X} " if m["is_error"] else ""
                print(f"{tag}  {err}{(m['content'] or '')[:200]}")
            else:
                print(f"{tag}  {(m['content'] or '')[:200]}")
        return 0
    if args.run_id:
        r = db.get_run(args.run_id)
        if not r:
            print(f"{C.R}No run {args.run_id}.{C.X}")
            return 1
        print(f"{C.BOLD}Run {r['run_id']}{C.X}  [{r['kind']}] {r['name']}  -> {r['target']}")
        print(f"  status={r['status']} phase={r['phase']} cost=${r['total_cost'] or 0:.4f} "
              f"findings={r['n_findings']}" +
              (f" solved={'yes' if r['solved'] else 'no'} flag={r['flag']}"
               if r['benchmark_id'] else ""))
        print(f"  {C.BOLD}Agents:{C.X}")
        for a in db.list_agents(args.run_id):
            err = f" {C.R}ERROR{C.X}" if a["error"] else ""
            print(f"    {C.B}{a['agent_id'][:8]}{C.X} {a['label']:<22} {a['role']:<11} "
                  f"{a['status']:<6} {a['tool_calls']} calls  ${a['cost'] or 0:.4f}{err}")
        print(f"  {C.DIM}drill into an agent: pentest.py history --agent-id <id>{C.X}")
        return 0
    # default: list runs
    rows = db.list_runs(limit=args.limit, kind=args.kind or None)
    if not rows:
        print(f"{C.Y}No runs recorded yet. DB: {db.path}{C.X}")
        return 0
    print(f"{C.BOLD}{'run_id':<34}{'kind':<10}{'name':<26}{'status':<9}"
          f"{'find':>5}{'cost':>9}  result{C.X}")
    for r in rows:
        res = ""
        if r["benchmark_id"]:
            res = (f"{C.G}SOLVED{C.X} {r['flag'] or ''}" if r["solved"]
                   else f"{C.R}miss{C.X}")
        print(f"{r['run_id']:<34}{(r['kind'] or ''):<10}{(r['name'] or '')[:24]:<26}"
              f"{(r['status'] or ''):<9}{r['n_findings'] or 0:>5}"
              f"{('$%.2f' % (r['total_cost'] or 0)):>9}  {res}")
    print(f"\n{C.DIM}{len(rows)} run(s). Detail: pentest.py history --run-id <id>   "
          f"DB: {db.path}{C.X}")
    return 0


def cmd_scopes(_args) -> int:
    print(f"{C.BOLD}Available scopes:{C.X}")
    for key, s in ALL_SCOPES.items():
        print(f"  {C.B}{key:<14}{C.X} {s.title}")
    return 0


def _slug(s: str) -> str:
    return "".join(c if c.isalnum() else "-" for c in s.lower()).strip("-")[:40] or "engagement"


def main() -> int:
    p = argparse.ArgumentParser(description="pentestkit — multi-agent pentest framework")
    sub = p.add_subparsers(dest="cmd", required=True)

    pr = sub.add_parser("run", help="run a full engagement")
    pr.add_argument("--engagement", required=True)
    pr.add_argument("--out", default=None, help="output dir (default: output/<name>-<ts>)")
    pr.add_argument("--env", default=None, help="path to .env (default: ./.env)")
    pr.add_argument("--yes", action="store_true", help="skip the interactive auth confirm")
    pr.add_argument("--resume", action="store_true", help="resume from existing context.json")
    pr.add_argument("--skip-plan", action="store_true", help="skip the planning phase")
    pr.add_argument("--ui", action="store_true", help="serve the live graph UI during the run")
    pr.add_argument("--ui-port", type=int, default=8420, help="UI port (default 8420)")
    pr.add_argument("--run-kind", default="pentest", help="run kind tag for the DB "
                    "(pentest|fleet|benchmark)")
    pr.add_argument("--run-group", default="", help="group id linking runs (fleet/bench batch)")
    pr.add_argument("--benchmark-id", default="", help="benchmark id when kind=benchmark")
    pr.add_argument("--commit-worklog", action="store_true",
                    help="after the run, commit+push this run's worklog to the repo")

    pc = sub.add_parser("check", help="test whether a URL is in scope")
    pc.add_argument("--engagement", required=True)
    pc.add_argument("--url", required=True)

    ps = sub.add_parser("serve", help="serve the live UI for a run dir (live or finished)")
    ps.add_argument("--run-dir", required=True)
    ps.add_argument("--port", type=int, default=8420)

    pci = sub.add_parser("ci", help="headless run for CI: writes SARIF + exits non-zero on "
                         "severe findings")
    pci.add_argument("--engagement", required=True)
    pci.add_argument("--out", default=None)
    pci.add_argument("--env", default=None)
    pci.add_argument("--skip-plan", action="store_true")
    pci.add_argument("--fail-on", default="high",
                     choices=["critical", "high", "medium", "low", "none"],
                     help="exit non-zero if any confirmed finding is at/above this severity")

    pfin = sub.add_parser("finish", help="finish an interrupted run (verify/score/report on "
                          "existing context — no re-scanning)")
    pfin.add_argument("--engagement", required=True, help="the engagement YAML used for the run")
    pfin.add_argument("--run-dir", required=True, help="the run's output dir (has context.json)")
    pfin.add_argument("--env", default=None)
    pfin.add_argument("--skip-verify", action="store_true",
                      help="skip re-verifying leftover candidates; just score + report")

    pb = sub.add_parser("serve-bench", help="serve the benchmark scoreboard (follows latest run)")
    pb.add_argument("--results-dir", required=True,
                    help="benchmarks/xbow/results (root -> auto-follows latest run) "
                         "OR a specific results/<ts> dir")
    pb.add_argument("--port", type=int, default=8600)

    pcon = sub.add_parser("console",
                          help="unified UI: Fleet + Benchmarks + Traces tabs on one port")
    pcon.add_argument("--fleet-dir", default=None,
                      help="fleet dir to serve/launch into (default: output/console-<ts>)")
    pcon.add_argument("--results-dir", default=None,
                      help="benchmark results root (default: benchmarks/xbow/results)")
    pcon.add_argument("--port", type=int, default=8500)
    pcon.add_argument("--env", default=None)

    pfc = sub.add_parser("fleet-console",
                         help="interactive fleet dashboard — add & launch targets from the UI")
    pfc.add_argument("--dir", default=None,
                     help="fleet dir to serve/launch into (default: output/fleet-console-<ts>)")
    pfc.add_argument("--port", type=int, default=8500)
    pfc.add_argument("--env", default=None)

    pf = sub.add_parser("fleet", help="run many engagements concurrently + combined dashboard")
    pf.add_argument("--fleet", required=True, help="path to a fleet YAML")
    pf.add_argument("--out", default=None, help="fleet output dir (default: output/fleet-<ts>)")
    pf.add_argument("--env", default=None)
    pf.add_argument("--max-parallel", type=int, default=None,
                    help="override max targets in parallel")
    pf.add_argument("--dashboard-port", type=int, default=None)
    pf.add_argument("--yes", action="store_true", help="skip the interactive auth confirm")

    ph = sub.add_parser("history", help="browse recorded runs / agents / messages from the DB")
    ph.add_argument("--run-id", default=None, help="show one run's agents")
    ph.add_argument("--agent-id", default=None, help="dump one agent's full message transcript")
    ph.add_argument("--kind", default=None, help="filter run list by kind")
    ph.add_argument("--limit", type=int, default=50)
    ph.add_argument("--import", dest="import_disk", action="store_true",
                    help="backfill the DB from on-disk runs (output/ + benchmark results)")

    sub.add_parser("scopes", help="list available scope modules")

    args = p.parse_args()
    if args.cmd == "run":
        return asyncio.run(cmd_run(args))
    if args.cmd == "check":
        return cmd_check(args)
    if args.cmd == "serve":
        return cmd_serve(args)
    if args.cmd == "ci":
        return asyncio.run(cmd_ci(args))
    if args.cmd == "finish":
        return asyncio.run(cmd_finish(args))
    if args.cmd == "history":
        return cmd_history(args)
    if args.cmd == "console":
        return cmd_console(args)
    if args.cmd == "serve-bench":
        return cmd_serve_bench(args)
    if args.cmd == "fleet-console":
        return cmd_fleet_console(args)
    if args.cmd == "fleet":
        return asyncio.run(cmd_fleet(args))
    if args.cmd == "scopes":
        return cmd_scopes(args)
    return 1


if __name__ == "__main__":
    sys.exit(main())