#!/usr/bin/env python3
"""Prove one backup reader's OpenBao boundary on railiance01 before delivery admission.

Run as root on railiance01. Stdlib only and self-contained, so it can be piped:

    ssh railiance01 'sudo python3 - --application mas --expected-kv-version 1' \
        < scripts/backup_reader_boundary_probe.py

Positive: the dedicated ServiceAccount logs in with audience openbao, holds only
its one reader policy within the CCR bounds, and reads exactly the expected KV
version with exactly the two pair fields.
Negative: a wrong audience and a foreign role cannot log in; sibling data,
metadata, parent LIST, a KV write and an administration endpoint are denied.
The write probe sends CAS=0 against an existing path, so it can never replace
the pair even if a policy were wrong.

Prints one JSON receipt. Values, tokens and JWTs never leave process memory;
the receipt holds field names, HTTP statuses, versions and bounds only.
Exit 0 only when every check passes.
"""

from __future__ import annotations

import argparse
import base64
import datetime as dt
import ipaddress
import json
import re
import subprocess
import sys
import time
import urllib.error
import urllib.request

CLUSTER_UID = "a553c742-0115-43d4-99a4-a5ca56fe0786"
SA = "backup-primary-delivery"
FIELDS = ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"]
APPLICATIONS = ("coulomb-social", "email-connect", "mas")
KUBE_ENV = {
    "PATH": "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin",
    "KUBECONFIG": "/etc/rancher/k3s/k3s.yaml",
}


class ProbeError(Exception):
    """Fixed diagnostic codes only."""


def guard(condition, code):
    if not condition:
        raise ProbeError(code)


def policy(application):
    return f"workload-kv-read-{application}-backup-primary"


def role(application):
    return f"{application}-backup-primary-workload-kv-read"


def data_path(application):
    return f"platform/data/workloads/{application}/backup-primary"


class Native:
    """kubectl and OpenBao access on the node. HTTP errors return their status."""

    def __init__(self):
        self.address = None
        self.opener = urllib.request.build_opener(urllib.request.ProxyHandler({}), _NoRedirect())

    def kube(self, args):
        try:
            result = subprocess.run(
                ["kubectl", "--request-timeout=20s", *args],
                capture_output=True, check=False, timeout=30, env=KUBE_ENV,
            )
        except (OSError, subprocess.SubprocessError):
            raise ProbeError("kubernetes_command_failed") from None
        guard(result.returncode == 0 and len(result.stdout) <= 1048576, "kubernetes_command_failed")
        return result.stdout

    def prepare(self, application):
        ns = json.loads(self.kube(["get", "namespace", "kube-system", "-o", "json"]))
        guard(ns.get("metadata", {}).get("uid") == CLUSTER_UID, "cluster_identity_mismatch")
        svc = json.loads(self.kube(["-n", "openbao", "get", "service", "openbao", "-o", "json"]))
        addr = ipaddress.ip_address(svc.get("spec", {}).get("clusterIP", ""))
        guard(
            addr in ipaddress.ip_network("10.43.0.0/16")
            and any(p.get("port") == 8200 for p in svc["spec"].get("ports", [])),
            "openbao_service_mismatch",
        )
        self.address = f"http://{addr}:8200/v1/"
        sa = json.loads(self.kube(["-n", application, "get", "serviceaccount", SA, "-o", "json"]))
        guard(sa.get("automountServiceAccountToken") is False, "non_mounting_service_account_required")
        return sa.get("metadata", {}).get("uid")

    def jwt(self, application, audience):
        token = self.kube([
            "-n", application, "create", "token", SA, f"--audience={audience}", "--duration=10m",
        ]).decode().strip()
        guard(len(token.split(".")) == 3 and len(token) < 16384, "bounded_jwt_required")
        return token

    def http(self, method, path, payload=None, token=None):
        headers = {"Content-Type": "application/json"}
        if token is not None:
            headers["X-Vault-Token"] = token
        body = None if payload is None else json.dumps(payload).encode()
        request = urllib.request.Request(self.address + path, data=body, headers=headers, method=method)
        try:
            with self.opener.open(request, timeout=20) as response:
                raw = response.read(1048577)
                guard(len(raw) <= 1048576, "openbao_response_too_large")
                return response.status, ({} if not raw else json.loads(raw))
        except urllib.error.HTTPError as exc:
            return exc.code, {}
        except (urllib.error.URLError, OSError):
            raise ProbeError("openbao_request_failed") from None


class _NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, *args, **kwargs):
        return None


def jwt_claims(token):
    piece = token.split(".")[1]
    return json.loads(base64.urlsafe_b64decode(piece + "=" * (-len(piece) % 4)))


def _login(native, application, login_role, audience="openbao"):
    status, body = native.http(
        "POST", "auth/kubernetes/login", {"role": login_role, "jwt": native.jwt(application, audience)}
    )
    auth = body.get("auth") if isinstance(body, dict) else None
    token = auth.get("client_token") if isinstance(auth, dict) else None
    return status, auth or {}, token if isinstance(token, str) and token else None


def _revoke(native, token):
    status, _ = native.http("POST", "auth/token/revoke-self", {}, token=token)
    return status in (200, 204)


def _denials(native, token, probes):
    """Run up to three denied probes on one four-use token, then revoke it."""
    statuses = {}
    try:
        for name, method, path, payload in probes:
            statuses[name], _ = native.http(method, path, payload, token=token)
    finally:
        revoked = _revoke(native, token)
    return statuses, revoked


def probe(native, application, expected_version, sa_uid=None):
    guard(application in APPLICATIONS, "application_refused")
    guard(type(expected_version) is int and expected_version > 0, "expected_version_required")
    expected_policy = policy(application)
    receipt = {
        "schema": "platform.backup-reader-boundary.v1",
        "application": application,
        "cluster_uid": CLUSTER_UID,
        "service_account": f"{application}/{SA}",
        "role": role(application),
        "policy": expected_policy,
        "kv_path": data_path(application),
        "expected_kv_version": expected_version,
        "values_recorded": False,
    }
    if sa_uid is None:
        sa_uid = native.prepare(application)
    claims = jwt_claims(native.jwt(application, "openbao"))
    bound = claims.get("kubernetes.io", {}).get("serviceaccount", {})
    receipt["jwt_binding"] = (
        claims.get("sub") == f"system:serviceaccount:{application}:{SA}"
        and claims.get("aud") == ["openbao"]
        and bool(bound.get("uid"))
        and bound.get("uid") == sa_uid
        and 0 < claims.get("exp", 0) - time.time() <= 660
    )

    # Positive: one token for lookup, the read, and revoke.
    status, auth, token = _login(native, application, role(application))
    positive = {"login_status": status}
    if token:
        try:
            positive["login_policies"] = auth.get("policies")
            positive["lease_duration"] = auth.get("lease_duration")
            _, info = native.http("GET", "auth/token/lookup-self", token=token)
            data = info.get("data", {}) if isinstance(info, dict) else {}
            positive["token"] = {
                "policies": data.get("policies"),
                "identity_policies": data.get("identity_policies") or [],
                "ttl": data.get("ttl"),
                "explicit_max_ttl": data.get("explicit_max_ttl"),
                "period": data.get("period") or 0,
                "num_uses": data.get("num_uses"),
            }
            read_status, record = native.http("GET", data_path(application), token=token)
            payload = record.get("data", {}) if isinstance(record, dict) else {}
            metadata = payload.get("metadata", {}) if isinstance(payload, dict) else {}
            values = payload.get("data", {}) if isinstance(payload, dict) else {}
            values = values if isinstance(values, dict) else {}
            positive["read_status"] = read_status
            positive["kv_version"] = metadata.get("version")
            positive["kv_destroyed"] = metadata.get("destroyed")
            positive["kv_deleted"] = bool(metadata.get("deletion_time"))
            positive["field_names"] = sorted(values)
            positive["encoding_ok"] = all(
                isinstance(values.get(name), str)
                and re.fullmatch(r"[A-Za-z0-9/+=_.-]{8,256}", values[name]) is not None
                for name in FIELDS
            )
            positive["fields_distinct"] = values.get(FIELDS[0]) != values.get(FIELDS[1])
            del values, payload, record
        finally:
            positive["token_revoked"] = _revoke(native, token)
    token_view = positive.get("token", {})
    positive["passed"] = (
        status == 200
        and positive.get("login_policies") == [expected_policy]
        and 0 < (positive.get("lease_duration") or 0) <= 300
        and token_view.get("policies") == [expected_policy]
        and token_view.get("identity_policies") == []
        and 0 < (token_view.get("ttl") or 0) <= 300
        and 0 < (token_view.get("explicit_max_ttl") or 0) <= 300
        and token_view.get("period") == 0
        and 0 < (token_view.get("num_uses") or 0) <= 4
        and positive.get("read_status") == 200
        and positive.get("kv_version") == expected_version
        and positive.get("kv_destroyed") is False
        and positive.get("kv_deleted") is False
        and positive.get("field_names") == FIELDS
        and positive.get("encoding_ok") is True
        and positive.get("fields_distinct") is True
        and positive.get("token_revoked") is True
    )
    receipt["positive"] = positive

    negative = {}
    # Login boundary: a wrong audience, and this account against a sibling role.
    status, _, stray = _login(native, application, role(application), audience="not-openbao")
    negative["wrong_audience_login"] = status
    if stray:
        _revoke(native, stray)
    sibling = next(item for item in APPLICATIONS if item != application)
    status, _, stray = _login(native, application, role(sibling))
    negative["foreign_role_login"] = status
    if stray:
        _revoke(native, stray)

    others = [item for item in APPLICATIONS if item != application]
    status, _, token = _login(native, application, role(application))
    guard(status == 200 and token, "negative_login_failed")
    first, revoked_first = _denials(native, token, [
        ("sibling_data_read", "GET", data_path(others[0]), None),
        ("metadata_read", "GET", f"platform/metadata/workloads/{application}/backup-primary", None),
        ("parent_list", "LIST", f"platform/metadata/workloads/{application}", None),
    ])
    status, _, token = _login(native, application, role(application))
    guard(status == 200 and token, "negative_login_failed")
    second, revoked_second = _denials(native, token, [
        # CAS=0 on an existing path cannot replace the pair even if allowed.
        ("kv_write", "POST", data_path(application), {"options": {"cas": 0}, "data": {"probe": "x"}}),
        ("second_sibling_data_read", "GET", data_path(others[1]), None),
        ("admin_policy_list", "LIST", "sys/policies/acl", None),
    ])
    negative.update(first)
    negative.update(second)
    negative["tokens_revoked"] = revoked_first and revoked_second
    login_denied = all(negative[name] in (400, 403) for name in ("wrong_audience_login", "foreign_role_login"))
    api_denied = all(
        negative[name] == 403
        for name in (
            "sibling_data_read", "metadata_read", "parent_list",
            "kv_write", "second_sibling_data_read", "admin_policy_list",
        )
    )
    negative["passed"] = login_denied and api_denied and negative["tokens_revoked"]
    receipt["negative"] = negative
    receipt["passed"] = receipt["jwt_binding"] and positive["passed"] and negative["passed"]
    receipt["checked_at"] = dt.datetime.now(dt.timezone.utc).isoformat()
    return receipt


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--application", choices=APPLICATIONS, required=True)
    parser.add_argument("--expected-kv-version", type=int, required=True)
    args = parser.parse_args(argv)
    try:
        receipt = probe(Native(), args.application, args.expected_kv_version)
    except ProbeError as exc:
        receipt = {"status": "failed", "code": str(exc), "values_recorded": False}
    except Exception as exc:  # noqa: BLE001 - never print backend bodies
        receipt = {"status": "failed", "code": "contained_probe_failure",
                   "error_type": type(exc).__name__, "values_recorded": False}
    print(json.dumps(receipt, indent=2, sort_keys=True))
    return 0 if receipt.get("passed") is True else 1


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