#!/usr/bin/env python3
"""AICEP v2.2 ingest gate (proposed Phase 2 extension). The model never does crypto and never sees a hand paste.

1. Issue a fresh 128-bit nonce.
2. Run aicep_sub.py (pinned by sha256) with --nonce and --attest-out, or accept an attestation file (--attestation).
3. Check: verifier signature against the REGISTERED verifier pubkey; nonce match; ts fresh (<=300 s,
   not in the future beyond 30 s); script_sha256 == pinned == sha256 of the script on disk;
   manifest_sha256 == sha256 of manifest.json on disk NOW; every ciphertext on disk NOW matches the manifest
   (TOCTOU); folder holds only listed files; result_sha256 == hash of the embedded result;
   exit_code/status consistent (0 <-> accepted).
4. Only then print gate output {"gate":"pass", "attestation":..., "content": <untrusted result>}.
   Otherwise {"gate":"fail","reason":...} and exit 4. No content on failure.
Config (JSON, human-maintained): {"verifier_pub": "<b64 raw>", "pinned_script_sha256": "<hex>"}
"""
import argparse, json, os, secrets, subprocess, sys, time
from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
from aicep_common import canonical, key_id, sha256_hex, unb64

HERE = os.path.dirname(os.path.abspath(__file__))
SUB = os.path.join(HERE, "aicep_sub.py")
MAX_AGE, MAX_SKEW = 300, 30


class GateFail(Exception):
    pass


def check(bundle, nonce, batch_dir, cfg, now=None, sub_path=SUB):
    now = int(now if now is not None else time.time())
    if not isinstance(bundle, dict) or not {"attestation", "sig", "result"} <= set(bundle):
        raise GateFail("no attestation/signature (unsigned result)")
    att, res = bundle["attestation"], bundle["result"]
    pub = unb64(cfg["verifier_pub"])
    if att.get("verifier_key_id") != key_id(pub):
        raise GateFail("attestation not from the registered verifier key")
    try:
        Ed25519PublicKey.from_public_bytes(pub).verify(unb64(bundle["sig"]), canonical(att))
    except (InvalidSignature, ValueError):
        raise GateFail("bad verifier signature")
    if att.get("nonce") != nonce:
        raise GateFail("nonce mismatch (stale or replayed attestation)")
    ts = att.get("ts")
    if not isinstance(ts, int) or now - ts > MAX_AGE or ts - now > MAX_SKEW:
        raise GateFail("attestation not fresh")
    disk_script = sha256_hex(open(sub_path, "rb").read())
    if not (att.get("script_sha256") == cfg["pinned_script_sha256"] == disk_script):
        raise GateFail("verifier script hash mismatch (modified aicep_sub.py)")
    if att.get("result_sha256") != sha256_hex(canonical(res)):
        raise GateFail("result does not match attestation")
    if (att.get("exit_code") == 0) != (att.get("status") == "accepted") or res.get("status") != att.get("status"):
        raise GateFail("exit code/status inconsistent")
    mb = open(os.path.join(batch_dir, "manifest.json"), "rb").read()
    if att.get("manifest_sha256") != sha256_hex(mb):
        raise GateFail("manifest on disk differs from attested manifest")
    m = json.loads(mb)
    # v2.3 QoP floor: signed + encrypted only
    if m.get("alg") != "age-x25519+ed25519" or any(not f.get("id") for f in m.get("files", [])):
        raise GateFail("QoP floor violated (sign+encrypt only)")
    if att.get("batch_id") not in (None, m.get("batch_id")) or (att.get("status") == "accepted" and
            (att.get("batch_id"), att.get("seq")) != (m.get("batch_id"), m.get("seq"))):
        raise GateFail("attestation is for a different batch")
    expected = {f"{f['id']}.age" for f in m.get("files", [])} | {"manifest.json", "manifest.json.sig"}
    if set(os.listdir(batch_dir)) != expected:
        raise GateFail("batch folder changed after verify (TOCTOU)")
    for f in m.get("files", []):
        if sha256_hex(open(os.path.join(batch_dir, f"{f['id']}.age"), "rb").read()) != f["sha256_ct"]:
            raise GateFail("file changed after verify (TOCTOU)")
    if att.get("status") != "accepted":
        raise GateFail(f"verifier {att.get('status')}: {att.get('reason')}")
    return {"gate": "pass", "attestation": att, "content": res}


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("batch_dir")
    ap.add_argument("--config", required=True)
    ap.add_argument("--compartment", required=True)
    ap.add_argument("--registry", required=True)
    ap.add_argument("--state", required=True)
    ap.add_argument("--folder", required=True, help="v2.5: sender folder the batch came from")
    ap.add_argument("--ack-gap-to", metavar="SKID:SEQ", help="HUMAN ONLY: exact gap ack")
    ap.add_argument("--operator", help="named human operator (required with --ack-gap-to)")
    ap.add_argument("--attest-out", default=None)
    a = ap.parse_args()
    cfg = json.load(open(a.config))
    nonce = secrets.token_hex(16)
    out = a.attest_out or os.path.join("/dev/shm", f"aicep-att-{nonce}.json")
    cmd = [sys.executable, SUB, a.batch_dir, "--compartment", a.compartment, "--registry", a.registry,
           "--state", a.state, "--folder", a.folder, "--nonce", nonce, "--attest-out", out] + (["--ack-gap-to", a.ack_gap_to, "--operator", a.operator or ""] if a.ack_gap_to else [])
    subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    try:
        bundle = json.load(open(out))
        print(json.dumps(check(bundle, nonce, a.batch_dir, cfg), indent=2))
    except (GateFail, OSError, ValueError) as e:
        print(json.dumps({"gate": "fail", "reason": str(e)}))
        sys.exit(4)
    finally:
        if not a.attest_out and os.path.exists(out):
            os.remove(out)


if __name__ == "__main__":
    main()