main.py

from __future__ import annotations
import json
import uuid
from pathlib import Path

from models import Finding, Asset
from risk.scorer import RiskScorer
from agent.planner import DefenderPlanner
from policy.engine import PolicyEngine
from broker.executor import ExecutionBroker
from verifier.verify import Verifier
from audit.logger import AuditLogger

WIDTH = 64


def _bar(char: str = "─") -> str:
    return char * WIDTH


def _load_findings() -> list[Finding]:
    data = json.loads(Path("examples/findings.json").read_text())
    return [Finding(**f) for f in data]


def _load_assets() -> dict[str, Asset]:
    data = json.loads(Path("examples/assets.json").read_text())
    return {a["id"]: Asset(**a) for a in data}


def main() -> None:
    print(_bar("━"))
    print("  Autonomous Defense Policy Agent")
    print(_bar("━"))

    findings = _load_findings()
    assets = _load_assets()
    print(f"\n  {len(findings)} findings  ·  {len(assets)} assets\n")
    print(_bar())

    # Clear previous run's audit log
    Path("audit/audit_log.jsonl").unlink(missing_ok=True)

    scorer  = RiskScorer()
    planner = DefenderPlanner()
    policy  = PolicyEngine()
    broker  = ExecutionBroker()
    verifier = Verifier()
    audit   = AuditLogger()

    llm_mode = planner._llm is not None
    print(f"  planner mode: {'LLM (Claude)' if llm_mode else 'rule-based (set ANTHROPIC_API_KEY to enable LLM)'}\n")

    stats: dict[str, int] = {"auto": 0, "approval": 0, "denied": 0}
    review_queue: list[dict] = []

    for finding in findings:
        asset = assets.get(finding.asset_id)
        if not asset:
            print(f"[{finding.id}]  WARNING: asset {finding.asset_id} not found, skipping")
            continue

        print(f"\n[{finding.id}]  {finding.finding_type}")
        print(f"  asset:        {finding.asset}  ({finding.environment})")

        risk = scorer.score(finding, asset)
        print(f"  risk:         {risk.total}/10  {risk.label.upper()}")

        plan = planner.plan(finding, risk, asset)
        print(f"  plan:         {plan.action_type}")
        print(f"                {plan.action}")

        decision = policy.evaluate(plan)
        audit.log_decision(plan, decision)

        if decision.decision == "allow":
            print(f"  policy:       ✓  ALLOW — {decision.reason}")
            result = broker.execute(plan)
            audit.log_execution(result)
            verification = verifier.verify(result)
            audit.log_verification(finding.id, verification)
            exec_status = "✓  success" if result.success else "✗  failed"
            verify_status = "✓  verified" if verification.passed else "✗  failed — rollback triggered"
            print(f"  execution:    {exec_status}")
            print(f"  verification: {verify_status}")
            stats["auto"] += 1

        elif decision.decision == "requires_approval":
            queue_id = f"HR-{uuid.uuid4().hex[:4].upper()}"
            audit.log_human_review(plan, queue_id)
            print(f"  policy:       ⚠   REQUIRES APPROVAL — {decision.reason}")
            print(f"  queued:       {queue_id}")
            review_queue.append({"id": queue_id, "finding": finding.id, "action": plan.action_type})
            stats["approval"] += 1

        else:
            audit.log_denial(plan, decision)
            print(f"  policy:       ✗  DENY — {decision.reason}")
            print(f"  no action taken")
            stats["denied"] += 1

    # Summary
    audit_count = sum(1 for _ in open("audit/audit_log.jsonl"))
    total = len(findings)

    print(f"\n{_bar('━')}")
    print("  Summary")
    print(_bar("━"))
    print(f"  findings processed:   {total}")
    print(f"  auto-remediated:      {stats['auto']}")
    print(f"  queued for approval:  {stats['approval']}")
    print(f"  denied:               {stats['denied']}")

    if review_queue:
        print(f"\n  human review queue:")
        for item in review_queue:
            print(f"    [{item['id']}]  {item['finding']}  →  {item['action']}")

    print(f"\n  audit log: audit/audit_log.jsonl  ({audit_count} records)")
    print(_bar("━"))


if __name__ == "__main__":
    main()