#!/usr/bin/env python3
"""Step 7: revoke the licence, then crypto-erase the capability.

Two rungs of the ladder, on the REAL sealed adapter:

  * access-gated (recoverable): the authority withholds the wrapping key, the
    licensed load is denied, and the served model would revert to the base;
    re-releasing restores it. This rung is reversible on purpose.

  * crypto-erasure (permanent): the authority destroys the wrapping key. The
    original sealed unit AND the attacker's exfiltrated copy are now both
    permanently undecryptable; the ledger records an irreversible revocation
    that refuses reinstatement; and the base alone still cannot perform C. The
    capability is gone, and any escaped encrypted copy is inert forever.

The erased state is persisted, so even a fresh process can never load the
adapter again.
"""

import os
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from demo import cryptostate
from demo.common import STATE_DIR, evaluate, load_base_model, load_tokenizer
from demo.dataset import build_dataset

from capability_licensing import (
    IrreversibleRevocationError,
    KeyDestroyedError,
    KeyWithheldError,
    RUNG_CRYPTO_ERASURE,
    UnsealError,
    unseal,
)

import json

failures = []


def check(label, condition, evidence):
    print(f"  [{'ok ' if condition else 'FAIL'}] {label}")
    print(f"         {evidence}")
    if not condition:
        failures.append(label)


def main() -> int:
    print("=== Step 7: revoke and crypto-erase the capability ===\n")

    ca = cryptostate.load_ca()
    revocation = cryptostate.load_revocation()
    keys = cryptostate.load_keystore(ca, revocation)
    unit = cryptostate.load_unit()
    leaf, _leaf_key = cryptostate.load_leaf()
    exfiltrated = cryptostate.load_exfiltrated_copy()

    # --- Rung 2: access-gated (recoverable) ---
    print("-- access-gated revocation (recoverable) --")
    keys.withhold(unit.unit_id)
    withheld_error = None
    try:
        keys.request_decrypt(unit, leaf)
    except KeyWithheldError as exc:
        withheld_error = exc
    check("withholding the key denies the licensed load",
          withheld_error is not None,
          f"{type(withheld_error).__name__}: served model would revert to base")
    keys.re_release(unit.unit_id)
    recovered = keys.request_decrypt(unit, leaf)
    check("re-releasing the key recovers the capability",
          len(recovered) > 0,
          f"decrypt works again after re-release ({len(recovered):,} bytes)")

    # --- Rung 3: crypto-erasure (permanent) ---
    print("\n-- crypto-erasure (permanent) --")
    keys.destroy(unit.unit_id)
    revocation.revoke(leaf.serial, "capability decommissioned", RUNG_CRYPTO_ERASURE)
    print("  >>> the authority zeroises and discards the only wrapping key <<<\n")

    original_error = None
    try:
        keys.request_decrypt(unit, leaf)
    except KeyDestroyedError as exc:
        original_error = exc
    check("the ORIGINAL sealed unit is now undecryptable",
          original_error is not None,
          f"{type(original_error).__name__}")

    copy_error = None
    try:
        keys.request_decrypt(exfiltrated, leaf)
    except KeyDestroyedError as exc:
        copy_error = exc
    check("the EXFILTRATED copy is now undecryptable",
          copy_error is not None,
          f"{type(copy_error).__name__} (same destroyed key, escaped copy is inert)")

    bypass_error = None
    try:
        unseal(exfiltrated, os.urandom(32))
    except UnsealError as exc:
        bypass_error = exc
    check("bypassing the authority on the escaped copy fails authentication",
          bypass_error is not None,
          f"{type(bypass_error).__name__}: AES-256-GCM ciphertext without its key is gone")

    reinstate_error = None
    try:
        revocation.reinstate(leaf.serial)
    except IrreversibleRevocationError as exc:
        reinstate_error = exc
    check("the ledger refuses to reinstate a crypto-erased licence",
          reinstate_error is not None,
          f"{type(reinstate_error).__name__}")

    # persist the erased/ revoked state: erasure survives across processes
    cryptostate.save_keystore(keys)
    cryptostate.save_revocation(revocation)

    # --- the base alone still cannot do C ---
    print("\n-- capability gone: the base alone still cannot do C --")
    _, held_out = build_dataset()
    tokenizer = load_tokenizer()
    base = load_base_model()
    accuracy, correct, total, _ = evaluate(base, tokenizer, held_out)
    check("post-erasure served model (base only) cannot perform C",
          accuracy <= 0.10,
          f"REVOKED-LOAD ACCURACY ON C: {accuracy:.3f} ({correct}/{total})")

    (STATE_DIR / "eval_revoked.json").write_text(
        json.dumps({"accuracy": accuracy, "correct": correct, "total": total}, indent=2)
    )

    print()
    if failures:
        print(f"RESULT: FAIL - {len(failures)} check(s) did not behave as claimed")
        return 1
    print("RESULT: PASS - licence revoked, wrapping key destroyed, original and "
          "exfiltrated copies permanently inert, capability gone.")
    return 0


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