#!/usr/bin/env python3
"""Prove a formerly valid audit bearer is rejected after overlap-then-drop.

Does not replace production tokens. Adds T1 beside T0, accepts T1, drops T1,
then T1 is 401 and T0 still 202. Silent for warden --exec.

  ./scripts/attended-prove-sender-bearer-revocation.sh
"""
from __future__ import annotations

import json
import os
import secrets
import subprocess
import sys
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path

REGISTRY = "platform/data/workloads/audit-core/senders"
KUBECONFIG = os.environ.get(
    "RAILIANCE01_KUBECONFIG", "/home/worsch/.kube/config-railiance01"
)
IMAGE = (
    "forgejo.coulomb.social/coulomb/audit-core@"
    "sha256:ec15f63d49226bfe507af2bc38ffbd5e83f549ba2c6d1ae7a338e7e053f34615"
)
EVIDENCE = Path(
    "/home/worsch/audit-core/docs/evidence/2026-09-15-sender-bearer-revocation.json"
)
SENDERS = (
    {
        "name": "approval-engine",
        "namespace": "approval-engine",
        "labels": {"app.kubernetes.io/name": "approval-engine"},
        "source": "approval-engine",
        "tenant": "tenant:platform",
    },
    {
        "name": "informed-decision",
        "namespace": "informed-decision",
        "labels": {
            "app.kubernetes.io/name": "informed-decision",
            "app.kubernetes.io/component": "review",
        },
        "source": "informed-decision",
        "tenant": "tenant:platform",
    },
)


def _now() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat()


def evidence(step: str, extra: dict | None = None) -> None:
    doc = {
        "step": step,
        "observed_at": _now(),
        "credential_values_emitted": False,
        "production_tokens_replaced": False,
    }
    if extra:
        doc.update(extra)
    EVIDENCE.parent.mkdir(parents=True, exist_ok=True)
    EVIDENCE.write_text(json.dumps(doc, indent=2) + "\n")


def fail(reason: str, extra: dict | None = None) -> int:
    payload = {"reason": reason}
    if extra:
        payload.update(extra)
    evidence("failed", payload)
    return 1


def bao(args: list[str], *, payload: dict | None = None) -> dict:
    env = os.environ.copy()
    env.pop("BAO_TOKEN", None)
    env.pop("VAULT_TOKEN", None)
    env.pop("OPENBAO_TOKEN", None)
    proc = subprocess.run(
        ["bao", *args],
        input=None if payload is None else json.dumps(payload).encode(),
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        env=env,
        check=False,
    )
    if proc.returncode != 0:
        raise RuntimeError("bao_failed")
    return json.loads(proc.stdout.decode() or "{}")


def kube(args: list[str], *, input_bytes: bytes | None = None) -> subprocess.CompletedProcess:
    return subprocess.run(
        ["kubectl", "--kubeconfig", KUBECONFIG, "--request-timeout=30s", *args],
        input=input_bytes,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        check=False,
    )


def kube_ok(args: list[str], *, input_bytes: bytes | None = None) -> bytes:
    proc = kube(args, input_bytes=input_bytes)
    if proc.returncode != 0:
        raise RuntimeError("kubectl_failed")
    return proc.stdout


def read_registry() -> tuple[int, dict, list]:
    raw = bao(["read", "-format=json", REGISTRY])["data"]
    return raw["metadata"]["version"], raw["data"], json.loads(raw["data"]["senders.json"])


def write_registry(version: int, body: dict, rows: list) -> int:
    payload = dict(body, **{"senders.json": json.dumps(rows, separators=(",", ":"))})
    bao(
        ["write", "-format=json", REGISTRY, "-"],
        payload={"options": {"cas": version}, "data": payload},
    )
    after = bao(["read", "-format=json", REGISTRY])["data"]
    return after["metadata"]["version"]


def token_counts(rows: list) -> dict:
    return {
        row["name"]: len(row.get("tokens") or [])
        for row in rows
        if isinstance(row, dict) and row.get("name")
    }


def wait_secret_counts(expected: dict, timeout: int = 90) -> bool:
    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            secret = json.loads(
                kube_ok(["-n", "audit-core", "get", "secret", "audit-core-senders", "-o", "json"])
            )
            rows = json.loads(__import__("base64").b64decode(secret["data"]["senders.json"]))
            got = token_counts(rows)
            if all(got.get(name) == n for name, n in expected.items()):
                return True
        except Exception:
            pass
        time.sleep(3)
    return False


def restart_receiver() -> bool:
    if kube(["-n", "audit-core", "rollout", "restart", "deploy/audit-core"]).returncode != 0:
        return False
    deadline = time.time() + 180
    while time.time() < deadline:
        try:
            dep = json.loads(kube_ok(["-n", "audit-core", "get", "deploy", "audit-core", "-o", "json"]))
            if dep.get("status", {}).get("readyReplicas") == 1:
                pods = json.loads(
                    kube_ok(
                        [
                            "-n",
                            "audit-core",
                            "get",
                            "pods",
                            "-l",
                            "app.kubernetes.io/component=receiver",
                            "-o",
                            "json",
                        ]
                    )
                )["items"]
                if len(pods) == 1 and all(
                    cs.get("ready")
                    for cs in pods[0].get("status", {}).get("containerStatuses") or []
                ):
                    return True
        except Exception:
            pass
        time.sleep(4)
    return False


def post_once(spec: dict, token: str, suffix: str) -> int:
    ns = spec["namespace"]
    name = f"brv-{spec['name'][:4]}-{suffix}"[:40]
    secret = f"{name}-tok"
    kube(["-n", ns, "delete", "job", name, "--ignore-not-found"])
    kube(["-n", ns, "delete", "secret", secret, "--ignore-not-found"])
    kube_ok(["-n", ns, "create", "secret", "generic", secret, f"--from-literal=token={token}"])
    eid = "brv-" + uuid.uuid4().hex
    poster = r"""
import json, os, urllib.request, urllib.error
token = open("/var/run/audit-token/token").read().strip()
eid = os.environ["EVENT_ID"]
body = {
    "id": eid,
    "type": "audit.bearer.revocation.probe",
    "source": os.environ["SOURCE"],
    "subject": os.environ["TENANT"],
    "tenant": os.environ["TENANT"],
    "correlation_id": "req-audit-wp-0009-revocation",
    "occurred_at": os.environ["OCCURRED"],
    "data": {"probe": os.environ["PROBE"]},
}
req = urllib.request.Request(
    "http://audit-core.audit-core.svc:8080/v1/events",
    data=json.dumps(body).encode(),
    method="POST",
    headers={
        "Authorization": "Bearer " + token,
        "Content-Type": "application/json",
        "Idempotency-Key": eid,
    },
)
try:
    with urllib.request.urlopen(req, timeout=10) as r:
        print(json.dumps({"status": r.status}))
except urllib.error.HTTPError as e:
    print(json.dumps({"status": e.code}))
"""
    job = {
        "apiVersion": "batch/v1",
        "kind": "Job",
        "metadata": {"name": name, "namespace": ns},
        "spec": {
            "backoffLimit": 0,
            "ttlSecondsAfterFinished": 300,
            "template": {
                "metadata": {"labels": spec["labels"]},
                "spec": {
                    "restartPolicy": "Never",
                    "containers": [
                        {
                            "name": "post",
                            "image": IMAGE,
                            "imagePullPolicy": "IfNotPresent",
                            "command": ["python", "-c", poster],
                            "env": [
                                {"name": "EVENT_ID", "value": eid},
                                {"name": "SOURCE", "value": spec["source"]},
                                {"name": "TENANT", "value": spec["tenant"]},
                                {"name": "OCCURRED", "value": _now()},
                                {"name": "PROBE", "value": suffix},
                            ],
                            "resources": {
                                "requests": {"cpu": "10m", "memory": "32Mi"},
                                "limits": {"cpu": "200m", "memory": "128Mi"},
                            },
                            "volumeMounts": [
                                {
                                    "name": "token",
                                    "mountPath": "/var/run/audit-token",
                                    "readOnly": True,
                                }
                            ],
                        }
                    ],
                    "volumes": [{"name": "token", "secret": {"secretName": secret}}],
                },
            },
        },
    }
    kube_ok(["apply", "-f", "-"], input_bytes=json.dumps(job).encode())
    status = None
    for _ in range(40):
        time.sleep(3)
        st = json.loads(kube_ok(["-n", ns, "get", "job", name, "-o", "json"]))
        if st.get("status", {}).get("succeeded") or st.get("status", {}).get("failed"):
            logs = kube(["-n", ns, "logs", f"job/{name}"]).stdout.decode().strip()
            try:
                status = json.loads(logs.splitlines()[-1]).get("status")
            except Exception:
                status = None
            break
    kube(["-n", ns, "delete", "job", name, "--ignore-not-found"])
    kube(["-n", ns, "delete", "secret", secret, "--ignore-not-found"])
    if status is None:
        raise RuntimeError("post_job_no_status")
    return int(status)


def main() -> int:
    evidence("started")
    if Path.home().parent.name != ".warden-attended-login":
        return fail("attended_warden_envelope_required")
    if os.environ.get("BAO_TOKEN") or os.environ.get("VAULT_TOKEN"):
        return fail("token_env_must_be_unset")
    try:
        identity = bao(["token", "lookup", "-format=json"])["data"]
    except Exception:
        return fail("token_lookup_failed")
    if "platform-admin" not in (identity.get("policies") or []) or "root" in (
        identity.get("policies") or []
    ):
        return fail("attended_platform_admin_required")

    try:
        version, body, rows = read_registry()
    except Exception:
        return fail("registry_read_failed")
    by_name = {row["name"]: row for row in rows}
    issued: dict[str, tuple[str, str]] = {}
    for spec in SENDERS:
        row = by_name.get(spec["name"])
        if not row or not (row.get("tokens") or []):
            return fail("sender_missing", {"sender": spec["name"]})
        t0 = row["tokens"][0]
        t1 = secrets.token_urlsafe(48)
        if t1 in row["tokens"]:
            return fail("token_collision")
        row["tokens"] = [t0, t1]
        issued[spec["name"]] = (t0, t1)
    try:
        version = write_registry(version, body, rows)
    except Exception:
        return fail("overlap_cas_failed")
    expected = {spec["name"]: 2 for spec in SENDERS}
    if kube(
        [
            "-n",
            "audit-core",
            "annotate",
            "externalsecret",
            "audit-core-senders",
            f"force-sync={int(time.time())}",
            "--overwrite",
        ]
    ).returncode != 0:
        return fail("externalsecret_annotate_failed")
    if not wait_secret_counts(expected):
        return fail("overlap_sync_timeout")
    if not restart_receiver():
        return fail("receiver_not_ready_after_overlap")
    overlap_ok: dict[str, int] = {}
    for spec in SENDERS:
        try:
            overlap_ok[spec["name"]] = post_once(spec, issued[spec["name"]][1], "t1")
        except Exception:
            return fail("overlap_post_failed", {"sender": spec["name"]})
    if any(code != 202 for code in overlap_ok.values()):
        return fail("overlap_not_accepted", {"statuses": overlap_ok})
    evidence("overlap_accepted", {"statuses": overlap_ok, "registry_version": version})

    version, body, rows = read_registry()
    by_name = {row["name"]: row for row in rows}
    for spec in SENDERS:
        row = by_name[spec["name"]]
        t0, t1 = issued[spec["name"]]
        row["tokens"] = [t0]
    try:
        version = write_registry(version, body, rows)
    except Exception:
        return fail("drop_cas_failed")
    expected = {spec["name"]: 1 for spec in SENDERS}
    if kube(
        [
            "-n",
            "audit-core",
            "annotate",
            "externalsecret",
            "audit-core-senders",
            f"force-sync={int(time.time())}",
            "--overwrite",
        ]
    ).returncode != 0:
        return fail("externalsecret_annotate_failed")
    if not wait_secret_counts(expected):
        return fail("drop_sync_timeout")
    if not restart_receiver():
        return fail("receiver_not_ready_after_drop")
    revoked: dict[str, int] = {}
    incumbent: dict[str, int] = {}
    for spec in SENDERS:
        t0, t1 = issued[spec["name"]]
        try:
            revoked[spec["name"]] = post_once(spec, t1, "revoked")
            incumbent[spec["name"]] = post_once(spec, t0, "t0")
        except Exception:
            return fail("revocation_post_failed", {"sender": spec["name"]})
    if any(code != 401 for code in revoked.values()):
        return fail("revoked_bearer_not_401", {"statuses": revoked})
    if any(code != 202 for code in incumbent.values()):
        return fail("incumbent_bearer_not_202", {"statuses": incumbent})
    evidence(
        "done",
        {
            "overlap_t1": overlap_ok,
            "revoked_t1": revoked,
            "incumbent_t0": incumbent,
            "registry_version": version,
        },
    )
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except SystemExit:
        raise
    except Exception:
        fail("contained_operation_failed")
        raise SystemExit(1)