aios/scripts/control_plane/aios_control_plane_audit.py

#!/usr/bin/env python3
from __future__ import annotations
from pathlib import Path
import argparse, json, sys, time, re
from aios_control_plane_consistency import *

def jload(value):
    try:
        return json.loads(value or "{}")
    except Exception:
        return {}

def rows(path, query, params=()):
    try:
        return sqlite_rows(path, query, params)
    except Exception as exc:
        return [{"_error": str(exc)}]

def main() -> int:
    ap = argparse.ArgumentParser(description="AIOS read-only control-plane consistency auditor")
    ap.add_argument("--task-id", required=True)
    ap.add_argument("--admission-db", type=Path, default=Path("/var/lib/aios-gitops/bridge-v2-staging/task-admission.sqlite3"))
    ap.add_argument("--task-db", type=Path, default=Path("/var/lib/aios-agent-task/state.sqlite3"))
    ap.add_argument("--git-dir", type=Path, default=Path("/data/aios-gitops/repo.git"))
    ap.add_argument("--branch")
    ap.add_argument("--openapi", type=Path)
    ap.add_argument("--server-source", type=Path)
    ap.add_argument("--artifact-root", type=Path)
    ap.add_argument("--now", type=int)
    ap.add_argument("--out", type=Path, default=Path("aios-control-plane-audit.json"))
    ns = ap.parse_args()
    now = ns.now or int(time.time())
    findings = []
    admission = None
    task = None
    workers = []
    lease = None
    git = None
    evidence = EvidenceSnapshot()

    # Admission plane.
    try:
        receipt = sqlite_rows(ns.admission_db,
            "SELECT * FROM task_admission_receipts WHERE task_id=? ORDER BY rowid DESC LIMIT 1", (ns.task_id,))
        if receipt:
            r = dict(receipt[0])
            envelope_id = r.get("envelope_id")
            env = sqlite_rows(ns.admission_db,
                "SELECT * FROM task_admission_envelopes WHERE envelope_id=? LIMIT 1", (envelope_id,))
            grant = sqlite_rows(ns.admission_db,
                "SELECT * FROM task_admission_grants WHERE envelope_id=? ORDER BY rowid DESC LIMIT 1", (envelope_id,))
            ed = dict(env[0]) if env else {}
            gd = dict(grant[0]) if grant else {}
            gj = jload(gd.get("grant_json"))
            ej = jload(ed.get("envelope_json"))
            admission = AdmissionSnapshot(
                task_id=ns.task_id,
                envelope_id=envelope_id,
                admission_status=ed.get("status"),
                grant_id=gd.get("grant_id"),
                grant_expires_at=gj.get("expires_at"),
                one_time_use=gj.get("one_time_use"),
                authority_base_oid=r.get("authority_base_oid") or ej.get("authority_base_oid"),
                execution_digest=r.get("execution_digest") or ej.get("execution_digest"),
                requested_capabilities=tuple(ej.get("requested_capabilities") or ()),
                file_allowlist=tuple(ej.get("file_allowlist") or ()),
                admitted_at=r.get("created_at") or r.get("admitted_at"),
            )
    except Exception as exc:
        findings.append(Finding("ADMISSION_PLANE_UNREADABLE", Severity.BLOCKER, str(exc)))

    # TaskPlane.
    try:
        trows = sqlite_rows(ns.task_db, "SELECT * FROM tasks WHERE task_id=? LIMIT 1", (ns.task_id,))
        if trows:
            td = dict(trows[0])
            meta = jload(td.get("metadata_json"))
            required = jload(td.get("required_capabilities_json"))
            if isinstance(required, dict):
                required = required.get("capabilities") or []
            task = TaskSnapshot(
                task_id=ns.task_id,
                status=td.get("status"),
                claimed_by=td.get("claimed_by"),
                lease_expires_at=td.get("lease_expires_at"),
                attempt=td.get("attempt"),
                revision=td.get("revision"),
                created_at=td.get("created_at"),
                updated_at=td.get("updated_at"),
                deadline=td.get("deadline"),
                required_capabilities=tuple(required or ()),
                metadata=meta if isinstance(meta, dict) else {},
            )
        arows = sqlite_rows(ns.task_db, "SELECT * FROM agents")
        for row in arows:
            d = dict(row)
            caps = jload(d.get("capabilities_json"))
            if isinstance(caps, dict):
                caps = caps.get("capabilities") or []
            last = parse_iso_epoch(d.get("last_seen_at"))
            workers.append(WorkerSnapshot(
                worker_id=d.get("agent_id") or d.get("worker_id") or "unknown",
                status=d.get("status"),
                last_seen_epoch=last,
                capabilities=tuple(caps or ()),
            ))
    except Exception as exc:
        findings.append(Finding("TASK_PLANE_UNREADABLE", Severity.BLOCKER, str(exc)))

    branch = ns.branch
    if not branch and task and isinstance(task.metadata, dict):
        plan = task.metadata.get("execution_plan") or task.metadata.get("plan") or task.metadata
        if isinstance(plan, dict):
            branch = plan.get("target_branch")
    if branch:
        current = git_ref(ns.git_dir, f"refs/heads/{branch}")
        git = GitSnapshot(branch=branch, current_oid=current, expected_oid=admission.authority_base_oid if admission else None)
        if admission and current and admission.authority_base_oid and current != admission.authority_base_oid:
            findings.append(Finding(
                "BRANCH_BINDING_DRIFT",
                Severity.BLOCKER,
                "Current branch OID differs from Admission authority_base_oid.",
                {"current": current, "bound": admission.authority_base_oid},
                "Replan/rebind formally; do not reuse old execution_digest.",
            ))

    truth, derived = derive_truth(admission, task, workers, lease, git, evidence, now)
    findings.extend(derived)
    if task:
        findings.extend(check_capability_aliases(task.required_capabilities, workers))

    # Plan/artifact integrity if the immutable plan points to files.
    if task and ns.artifact_root:
        plan = task.metadata.get("execution_plan") or task.metadata.get("plan") or task.metadata
        if isinstance(plan, dict):
            hashes = plan.get("file_sha256") or {}
            for rel in plan.get("file_allowlist") or []:
                findings.extend(inspect_artifact(ns.artifact_root / rel, hashes.get(rel)))

    if ns.openapi and ns.server_source:
        findings.extend(check_openapi_backend_parity(
            ns.openapi, ns.server_source,
            required_ops=("getTaskAdmissionEnvelope","submitTaskAdmissionAuthorization","getTaskAdmissionStatus"),
        ))

    report = {
        "schema": "aios.control-plane-consistency-audit.v1",
        "observed_at": now,
        "task_id": ns.task_id,
        "truth_state": truth.value,
        "admission": asdict(admission) if admission else None,
        "task": asdict(task) if task else None,
        "eligible_workers": [asdict(w) for w in eligible_workers(task, workers, now)] if task else [],
        "git": asdict(git) if git else None,
        "findings": [asdict(f) for f in findings],
        "summary": {
            "blockers": sum(f.severity == Severity.BLOCKER for f in findings),
            "majors": sum(f.severity == Severity.MAJOR for f in findings),
            "minors": sum(f.severity == Severity.MINOR for f in findings),
        },
    }
    dump_report(ns.out, report)
    print(json.dumps(report["summary"], sort_keys=True))
    return 2 if report["summary"]["blockers"] else 1 if report["summary"]["majors"] else 0

if __name__ == "__main__":
    raise SystemExit(main())