#!/usr/bin/env python3
"""Step 5: package the LoRA as an encrypted, licensed capability unit.

The trained adapter (config + safetensors) is packed into one blob and
envelope-encrypted with the repository's clean-room crypto: a fresh AES-256-GCM
data key encrypts the adapter, and that data key is wrapped by a per-unit
wrapping key held only by the key authority. A three-tier Ed25519 certificate
chain is issued and the unit is licensed under the leaf capability certificate.

The sealed unit is pure ciphertext: safe to copy or exfiltrate, useless without
the wrapping key. An attacker's byte-for-byte copy is written out too, to be
used in step 7. All authority state is persisted to the gitignored state dir so
serve (06) and revoke (07) can run as separate commands.
"""

import hashlib
import sys
from pathlib import Path

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

from demo import cryptostate
from demo.common import ADAPTER_DIR, BASE_MODEL_ID, tar_dir_bytes

from capability_licensing import CertificateAuthority, KeyAuthority, RevocationList

UNIT_ID = "sigil-control-adapter-001"


def main() -> int:
    print("=== Step 5: package the adapter as an encrypted licensed unit ===\n")
    if not (ADAPTER_DIR / "adapter_config.json").exists():
        print(f"  no adapter at {ADAPTER_DIR}; run step 3 first.")
        return 1

    adapter_blob = tar_dir_bytes(ADAPTER_DIR)
    digest = hashlib.sha256(adapter_blob).hexdigest()
    print(f"  adapter blob:      {len(adapter_blob):,} bytes plaintext "
          f"(tar of {[p.name for p in sorted(ADAPTER_DIR.iterdir()) if p.is_file()]})")
    print(f"  adapter sha256:    {digest[:16]}...\n")

    # --- issue the certificate chain ---
    ca = CertificateAuthority()
    root = ca.create_root("Reference Root Authority")
    organisation = ca.issue_organisation(root.serial, "Reference Research Organisation")
    leaf, leaf_key = ca.issue_capability(
        organisation.serial,
        "sigil-control-protocol",
        claims={
            "capability": "SIGIL home-automation control directive protocol",
            "base_model": BASE_MODEL_ID,
            "grant": "research demonstration only",
        },
    )
    print("  certificate chain issued (Ed25519):")
    print(f"    root:         {root.serial}")
    print(f"    organisation: {organisation.serial}")
    print(f"    capability:   {leaf.serial}  <- the licence for this unit\n")

    # --- envelope-encrypt the adapter under the key authority ---
    revocation = RevocationList()
    keys = KeyAuthority(
        trusted_roots=[root],
        intermediates=[organisation],
        revocation_list=revocation,
    )
    sealed = keys.seal_unit(UNIT_ID, adapter_blob, leaf.serial)
    print(f"  sealed unit '{sealed.unit_id}':")
    print(f"    ciphertext:   {sealed.ciphertext_size():,} bytes AES-256-GCM")
    print(f"    licensed to:  {sealed.capability_serial}")
    print("    wrapping key: held ONLY at the key authority (never in the unit)\n")

    # --- persist all state for the next commands ---
    cryptostate.save_ca(ca)
    cryptostate.save_revocation(revocation)
    cryptostate.save_keystore(keys)
    cryptostate.save_unit(sealed)
    cryptostate.save_leaf(leaf, leaf_key)
    cryptostate.save_meta(root.serial, organisation.serial, leaf.serial, UNIT_ID)

    # --- an attacker steals a byte-for-byte copy of the ciphertext ---
    cryptostate.save_exfiltrated_copy(sealed)
    print("  an attacker exfiltrates a byte-for-byte copy of the sealed unit")
    print("  (saved to state/exfiltrated_unit.json -- pure ciphertext, no key)\n")

    print("RESULT: PASS - adapter sealed as an encrypted licensed capability unit.")
    return 0


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