#!/usr/bin/env python3
"""
Reference envelope crypto for the sealed-object contract (Tool 2).

seal(): plaintext -> (sealed-object metadata, blob bytes). A random per-object DEK
(AES-256-GCM) encrypts the content; the DEK is wrapped to each recipient's X25519
public key with a libsodium sealed box. P0 stores the plaintext as the blob with no
keys at all. Every object seal() produces passes sealed_object.validate().

unseal(): with a recipient private key, verify integrity, unwrap the DEK, decrypt.

Self-test: python3 envelope.py --selftest
"""

from __future__ import annotations

import base64
import hashlib
import os
import pathlib
import sys
import tempfile

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from sealed_object import validate, SCHEMA_VERSION  # noqa: E402

from cryptography.hazmat.primitives.ciphers.aead import AESGCM  # noqa: E402
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes  # noqa: E402
from cryptography.exceptions import InvalidTag  # noqa: E402
from nacl.public import PrivateKey, PublicKey, SealedBox  # noqa: E402


DEFAULT_STREAM_CHUNK_BYTES = 1024 * 1024


def _b64(b: bytes) -> str:
    return base64.b64encode(b).decode("ascii")


def _unb64(s: str) -> bytes:
    return base64.b64decode(s)


def gen_recipient_keypair():
    """Return (private_key_b64, public_key_b64) for a recipient (X25519)."""
    sk = PrivateKey.generate()
    return _b64(bytes(sk)), _b64(bytes(sk.public_key))


def public_key_from_private(private_b64: str) -> str:
    """Derive the base64 X25519 public key for a base64 private key."""
    return _b64(bytes(PrivateKey(_unb64(private_b64)).public_key))


def fingerprint(public_b64: str) -> str:
    return "x25519:" + hashlib.sha256(_unb64(public_b64)).hexdigest()[:32]


def seal(plaintext: bytes, *, posture: str, storage_class: str, residency_scope: str,
         obj_id: str, recipients=None, metadata=None):
    recipients = recipients or []
    metadata = dict(metadata or {})
    metadata.setdefault("size", len(plaintext))

    if posture == "P0":
        blob = plaintext
        alg = {"content": "none", "wrap": "none"}
        wrapped = []
    else:
        if not recipients:
            raise ValueError(f"{posture} requires at least one recipient public key")
        dek = AESGCM.generate_key(bit_length=256)
        nonce = os.urandom(12)
        ct = AESGCM(dek).encrypt(nonce, plaintext, None)
        blob = nonce + ct
        alg = {"content": "AES-256-GCM", "wrap": "x25519-sealedbox"}
        wrapped = [
            {
                "recipient_ref": fingerprint(pub),
                "wrapped_dek": _b64(SealedBox(PublicKey(_unb64(pub))).encrypt(dek)),
            }
            for pub in recipients
        ]

    integrity = {"ciphertext_sha256": hashlib.sha256(blob).hexdigest()}
    if posture == "P1":
        # Managed tier may keep a plaintext hash; zero-knowledge (P2) deliberately does not.
        integrity["plaintext_sha256"] = hashlib.sha256(plaintext).hexdigest()

    obj = {
        "version": SCHEMA_VERSION,
        "id": obj_id,
        "posture": posture,
        "storage_class": storage_class,
        "residency_scope": residency_scope,
        "alg": alg,
        "integrity": integrity,
        "payload_ref": f"blob/{obj_id}",
        "wrapped_deks": wrapped,
        "metadata": metadata,
    }
    errs = validate(obj)
    if errs:
        raise AssertionError("seal produced an invalid object: " + "; ".join(errs))
    return obj, blob


def seal_file(source_path, destination_path, *, posture: str, storage_class: str,
              residency_scope: str, obj_id: str, recipients=None, metadata=None,
              chunk_size: int = DEFAULT_STREAM_CHUNK_BYTES):
    """Seal a file without holding its plaintext or ciphertext in memory.

    The output keeps the byte format produced by :func:`seal`: P0 is the source
    bytes, and encrypted postures are ``nonce + ciphertext + GCM tag``. The
    returned tuple is ``(object, source_size, ciphertext_size, source_sha256)``.
    The caller owns the destination lifecycle and should remove a partial file
    if this function raises.
    """
    source = pathlib.Path(source_path)
    destination = pathlib.Path(destination_path)
    if source.resolve() == destination.resolve():
        raise ValueError("source and destination must be different files")
    chunk_size = max(1, int(chunk_size))
    if posture not in {"P0", "P1", "P2"}:
        raise ValueError("unsupported sealed-object posture")

    recipients = list(recipients or [])
    if posture != "P0" and not recipients:
        raise ValueError(f"{posture} requires at least one recipient public key")
    source_size = source.stat().st_size
    metadata = dict(metadata or {})
    metadata["size"] = source_size
    source_hash = hashlib.sha256()
    ciphertext_hash = hashlib.sha256()

    if posture == "P0":
        alg = {"content": "none", "wrap": "none"}
        wrapped = []
        with source.open("rb") as src, destination.open("wb") as dst:
            while True:
                chunk = src.read(chunk_size)
                if not chunk:
                    break
                source_hash.update(chunk)
                ciphertext_hash.update(chunk)
                dst.write(chunk)
        ciphertext_size = destination.stat().st_size
    else:
        dek = AESGCM.generate_key(bit_length=256)
        nonce = os.urandom(12)
        encryptor = Cipher(algorithms.AES(dek), modes.GCM(nonce)).encryptor()
        with source.open("rb") as src, destination.open("wb") as dst:
            dst.write(nonce)
            ciphertext_hash.update(nonce)
            while True:
                chunk = src.read(chunk_size)
                if not chunk:
                    break
                source_hash.update(chunk)
                encrypted = encryptor.update(chunk)
                if encrypted:
                    dst.write(encrypted)
                    ciphertext_hash.update(encrypted)
            encryptor.finalize()
            dst.write(encryptor.tag)
            ciphertext_hash.update(encryptor.tag)
        alg = {"content": "AES-256-GCM", "wrap": "x25519-sealedbox"}
        wrapped = [
            {
                "recipient_ref": fingerprint(pub),
                "wrapped_dek": _b64(SealedBox(PublicKey(_unb64(pub))).encrypt(dek)),
            }
            for pub in recipients
        ]
        ciphertext_size = destination.stat().st_size

    integrity = {"ciphertext_sha256": ciphertext_hash.hexdigest()}
    if posture == "P1":
        integrity["plaintext_sha256"] = source_hash.hexdigest()
    obj = {
        "version": SCHEMA_VERSION,
        "id": obj_id,
        "posture": posture,
        "storage_class": storage_class,
        "residency_scope": residency_scope,
        "alg": alg,
        "integrity": integrity,
        "payload_ref": f"blob/{obj_id}",
        "wrapped_deks": wrapped,
        "metadata": metadata,
    }
    errs = validate(obj)
    if errs:
        raise AssertionError("seal_file produced an invalid object: " + "; ".join(errs))
    return obj, source_size, ciphertext_size, source_hash.hexdigest()


def verify_integrity(obj, blob: bytes) -> bool:
    return hashlib.sha256(blob).hexdigest() == obj["integrity"]["ciphertext_sha256"]


def verify_integrity_file(obj, blob_path, chunk_size: int = DEFAULT_STREAM_CHUNK_BYTES) -> bool:
    """Verify a sealed ciphertext file without reading it into memory."""
    expected = str(obj.get("integrity", {}).get("ciphertext_sha256", ""))
    digest = hashlib.sha256()
    with pathlib.Path(blob_path).open("rb") as handle:
        while True:
            chunk = handle.read(max(1, int(chunk_size)))
            if not chunk:
                break
            digest.update(chunk)
    return digest.hexdigest() == expected


def _recipient_dek(obj, recipient_private_b64: str | None):
    if recipient_private_b64 is None:
        raise ValueError("a recipient private key is required to unseal P1/P2")
    sk = PrivateKey(_unb64(recipient_private_b64))
    my_ref = fingerprint(_b64(bytes(sk.public_key)))
    for wrapped in obj["wrapped_deks"]:
        if wrapped["recipient_ref"] == my_ref:
            return SealedBox(sk).decrypt(_unb64(wrapped["wrapped_dek"]))
    raise ValueError("no wrapped DEK addressed to this recipient")


def unseal(obj, blob: bytes, recipient_private_b64: str | None = None) -> bytes:
    if not verify_integrity(obj, blob):
        raise ValueError("integrity check failed: blob does not match ciphertext_sha256")
    if obj["posture"] == "P0":
        return blob
    dek = _recipient_dek(obj, recipient_private_b64)
    nonce, ct = blob[:12], blob[12:]
    return AESGCM(dek).decrypt(nonce, ct, None)


def unseal_file(obj, source_path, destination_path,
                recipient_private_b64: str | None = None,
                chunk_size: int = DEFAULT_STREAM_CHUNK_BYTES):
    """Verify, decrypt, and atomically write a sealed file with bounded memory.

    The source is hashed before decryption. Encrypted payloads are then decrypted in
    bounded chunks after separating the 12-byte nonce and final 16-byte GCM tag. The
    destination is written to a private sibling temporary file and replaced only after
    GCM authentication and declared plaintext-size/hash checks succeed. The returned tuple
    is ``(ciphertext_size, plaintext_size, plaintext_sha256)``.
    """
    source = pathlib.Path(source_path)
    destination = pathlib.Path(destination_path)
    if source.resolve() == destination.resolve():
        raise ValueError("source and destination must be different files")
    errors = validate(obj)
    if errors:
        raise ValueError("invalid sealed object: " + "; ".join(errors))
    if not source.is_file():
        raise FileNotFoundError(source)
    chunk_size = max(1, int(chunk_size))
    ciphertext_size = source.stat().st_size
    if not verify_integrity_file(obj, source, chunk_size):
        raise ValueError("integrity check failed: file does not match ciphertext_sha256")

    destination.parent.mkdir(parents=True, exist_ok=True)
    temporary = tempfile.NamedTemporaryFile(
        prefix=f".{destination.name}.", suffix=".silo-part", dir=str(destination.parent),
        delete=False)
    temporary_path = pathlib.Path(temporary.name)
    plaintext_hash = hashlib.sha256()
    plaintext_size = 0
    try:
        if obj["posture"] == "P0":
            with source.open("rb") as src:
                while True:
                    chunk = src.read(chunk_size)
                    if not chunk:
                        break
                    temporary.write(chunk)
                    plaintext_hash.update(chunk)
                    plaintext_size += len(chunk)
        else:
            if ciphertext_size < 12 + 16:
                raise ValueError("encrypted sealed file is shorter than nonce and tag")
            dek = _recipient_dek(obj, recipient_private_b64)
            with source.open("rb") as src:
                nonce = src.read(12)
                src.seek(ciphertext_size - 16)
                tag = src.read(16)
                ciphertext_bytes = ciphertext_size - 12 - 16
                src.seek(12)
                decryptor = Cipher(
                    algorithms.AES(dek), modes.GCM(nonce, tag)).decryptor()
                remaining = ciphertext_bytes
                while remaining:
                    chunk = src.read(min(chunk_size, remaining))
                    if not chunk or len(chunk) > remaining:
                        raise ValueError("encrypted sealed file ended before its declared tag")
                    remaining -= len(chunk)
                    plaintext = decryptor.update(chunk)
                    if plaintext:
                        temporary.write(plaintext)
                        plaintext_hash.update(plaintext)
                        plaintext_size += len(plaintext)
                try:
                    plaintext = decryptor.finalize()
                except InvalidTag as exc:
                    raise ValueError("sealed file authentication failed") from exc
                if plaintext:
                    temporary.write(plaintext)
                    plaintext_hash.update(plaintext)
                    plaintext_size += len(plaintext)
        temporary.flush()
        os.fsync(temporary.fileno())
        temporary.close()
        declared_size = obj.get("metadata", {}).get("size")
        if declared_size is not None and int(declared_size) != plaintext_size:
            raise ValueError("sealed file plaintext size does not match its metadata")
        declared_plaintext_hash = obj.get("integrity", {}).get("plaintext_sha256")
        actual_plaintext_hash = plaintext_hash.hexdigest()
        if declared_plaintext_hash and declared_plaintext_hash != actual_plaintext_hash:
            raise ValueError("sealed file plaintext hash does not match its envelope")
        os.replace(temporary_path, destination)
        return ciphertext_size, plaintext_size, actual_plaintext_hash
    except Exception:
        try:
            temporary.close()
        except Exception:
            pass
        try:
            temporary_path.unlink()
        except FileNotFoundError:
            pass
        raise


def add_recipient(obj, granter_private_b64: str, new_recipient_public_b64: str):
    """Grant read access: re-wrap the object's DEK to a new recipient's public key.

    The granter must currently be a recipient; they unwrap the DEK with their private key
    and wrap it to the new key. The content is NOT re-encrypted; only a new wrapped DEK is
    added. Returns the updated object. This is how a read permission is granted.
    """
    if obj["posture"] == "P0":
        raise ValueError("P0 objects carry no keys to grant")
    sk = PrivateKey(_unb64(granter_private_b64))
    my_ref = fingerprint(_b64(bytes(sk.public_key)))
    dek = None
    for w in obj["wrapped_deks"]:
        if w["recipient_ref"] == my_ref:
            dek = SealedBox(sk).decrypt(_unb64(w["wrapped_dek"]))
            break
    if dek is None:
        raise ValueError("granter is not a current recipient; cannot grant read")
    new_ref = fingerprint(new_recipient_public_b64)
    if any(w["recipient_ref"] == new_ref for w in obj["wrapped_deks"]):
        return obj  # already a recipient
    wrapped = _b64(SealedBox(PublicKey(_unb64(new_recipient_public_b64))).encrypt(dek))
    new_obj = dict(obj)
    new_obj["wrapped_deks"] = obj["wrapped_deks"] + [{"recipient_ref": new_ref, "wrapped_dek": wrapped}]
    return new_obj


# --------------------------------------------------------------------------- selftest

def _selftest() -> int:
    import tempfile

    passed = failed = 0

    def check(name, cond, detail=""):
        nonlocal passed, failed
        if cond:
            passed += 1
            print(f"  [PASS] {name}")
        else:
            failed += 1
            print(f"  [FAIL] {name}{('  ' + detail) if detail else ''}")

    msg = b"privileged matter file: contents that must never be readable by the host"

    # P0 plain round-trip.
    o0, b0 = seal(msg, posture="P0", storage_class="hot", residency_scope="QC", obj_id="o0")
    check("P0 round-trip", unseal(o0, b0) == msg)
    check("P0 produces valid object", validate(o0) == [])
    check("P0 blob is the plaintext (plain tier)", b0 == msg)

    # P2 single recipient.
    sk_a, pk_a = gen_recipient_keypair()
    check("private key derives its public key", public_key_from_private(sk_a) == pk_a)
    o2, b2 = seal(msg, posture="P2", storage_class="cold", residency_scope="CA",
                  obj_id="o2", recipients=[pk_a])
    check("P2 round-trip with correct key", unseal(o2, b2, sk_a) == msg)
    check("P2 produces valid object", validate(o2) == [])
    check("P2 blob is NOT plaintext (encrypted)", b2 != msg)

    # P2 multi-recipient: both can read.
    sk_b, pk_b = gen_recipient_keypair()
    om, bm = seal(msg, posture="P2", storage_class="hot", residency_scope="CA",
                  obj_id="om", recipients=[pk_a, pk_b])
    check("multi-recipient: A can read", unseal(om, bm, sk_a) == msg)
    check("multi-recipient: B can read", unseal(om, bm, sk_b) == msg)

    # Wrong recipient cannot read.
    sk_c, _ = gen_recipient_keypair()
    try:
        unseal(o2, b2, sk_c)
        check("non-recipient is refused", False, "decrypt unexpectedly succeeded")
    except ValueError:
        check("non-recipient is refused", True)

    # Tamper detection.
    tampered = bytearray(b2)
    tampered[-1] ^= 0x01
    try:
        unseal(o2, bytes(tampered), sk_a)
        check("tamper is detected", False, "no error on modified blob")
    except ValueError:
        check("tamper is detected", True)

    # Granting read: A (a recipient of o2) grants a new key G; G can then read; B cannot grant.
    sk_g, pk_g = gen_recipient_keypair()
    o2g = add_recipient(o2, sk_a, pk_g)
    check("grant: new key can read after A grants", unseal(o2g, b2, sk_g) == msg)
    try:
        add_recipient(o2, sk_b, pk_g)  # B is not a recipient of o2, so cannot grant
        check("non-recipient cannot grant read", False)
    except ValueError:
        check("non-recipient cannot grant read", True)

    with tempfile.TemporaryDirectory(prefix="havre-envelope-stream-") as tmp:
        source = os.path.join(tmp, "source.bin")
        sealed = os.path.join(tmp, "sealed.bin")
        with open(source, "wb") as handle:
            handle.write(msg * 4097)
        file_obj, source_size, ciphertext_size, source_hash = seal_file(
            source, sealed, posture="P2", storage_class="hot", residency_scope="QC",
            obj_id="streamed", recipients=[pk_a], chunk_size=17)
        with open(sealed, "rb") as handle:
            file_blob = handle.read()
        with open(source, "rb") as handle:
            source_bytes = handle.read()
        check("file sealing reports bounded-stream sizes",
              source_size == len(source_bytes) and ciphertext_size == len(file_blob))
        check("file sealing reports the source hash",
              source_hash == hashlib.sha256(source_bytes).hexdigest())
        check("file sealing keeps the envelope compatible",
              validate(file_obj) == [] and unseal(file_obj, file_blob, sk_a) == source_bytes)
        restored = os.path.join(tmp, "restored.bin")
        restored_size = unseal_file(file_obj, sealed, restored, sk_a, chunk_size=17)
        check("file unsealing writes the original bytes",
              restored_size[1] == len(source_bytes)
              and open(restored, "rb").read() == source_bytes)
        tampered_path = os.path.join(tmp, "tampered.bin")
        tampered_blob = bytearray(file_blob)
        tampered_blob[-1] ^= 0x01
        with open(tampered_path, "wb") as handle:
            handle.write(tampered_blob)
        with open(restored, "wb") as handle:
            handle.write(b"keep-existing-target")
        try:
            unseal_file(file_obj, tampered_path, restored, sk_a, chunk_size=17)
            check("file unsealing refuses ciphertext hash mismatch", False)
        except ValueError:
            check("file unsealing refuses ciphertext hash mismatch",
                  open(restored, "rb").read() == b"keep-existing-target")

    print(f"\n{passed} passed, {failed} failed")
    return 1 if failed else 0


if __name__ == "__main__":
    if "--selftest" in sys.argv:
        print(f"envelope.py self-test ({SCHEMA_VERSION})\n")
        raise SystemExit(_selftest())
    print(__doc__)