#!/usr/bin/env python3
"""
Hybrid encrypt a file using AES-256-GCM + ML-KEM-1024.
Streaming handles files larger than RAM.
Usage:
    python3 encrypt.py <filename>

Output: <filename>.kyb
Requires: mlkem1024.pub in the same directory.
"""

import os
import sys
import cbor2
import oqs
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

MLKEM_KEM    = "ML-KEM-1024"
SIGNATURE    = b"KYB"       # 3-byte file signature
VERSION      = 1            # 1-byte format version
KDF_NAME     = "HKDF-SHA256"      # KDF advertised in the header
KDF_INFO     = "mlkem-aes-key"    # HKDF domain-separation label
OUT_SUFFIX   = ".kyb"
IV_LEN       = 12   # GCM standard nonce length (96 bits)
TAG_LEN      = 16   # full GCM auth tag (128 bits)
CHUNK_SIZE   = 64 * 1024 * 1024   # 64 MB


def build_header(mlkem_ct_len: int, iv_len: int, tag_len: int) -> bytes:
    """Build the CBOR header: signature + version + CBOR map of fields.

    CBOR (RFC 8949) is compact, self-delimiting, and standardized in IETF
    crypto formats (e.g. COSE), so the header stays forward-compatible.
    """
    header_map = {
        "KEM": "MLKEM1024",
        "CIPHER": "AES256GCM",
        "MLKEMCT": mlkem_ct_len,
        "IV": iv_len,
        "TAG": tag_len,
        "KDF": "HKDF-SHA256",
        "INFO": "mlkem-aes-key"
    }
    return SIGNATURE + bytes([VERSION]) + cbor2.dumps(header_map, canonical=True)


def derive_key(shared_secret: bytes, salt: bytes) -> bytes:
    """Derive the single AES-256-GCM key from the ML-KEM shared secret.

    The nonce is used as the HKDF salt, binding the key to this specific
    encryption (defense-in-depth: key/nonce mismatches fail loudly).
    """
    return HKDF(
        algorithm=hashes.SHA256(), length=32, salt=salt,
        info=KDF_INFO.encode("ascii")
    ).derive(shared_secret)


def encrypt_file(plaintext_path: str, output_path: str, pub_key: bytes):
    iv = os.urandom(IV_LEN)

    # Each file gets a brand-new AES key:
    #   fresh shared secret (encap_secret) -> fresh derived key (derive_key).
    # Since every file has its own unique key, even a repeated nonce between
    # files would be safe. Never cache or reuse a shared secret across files -
    # doing so would break that guarantee.
    kem = oqs.KeyEncapsulation(MLKEM_KEM)
    mlkem_ct, shared_secret = kem.encap_secret(pub_key)
    kem.free()

    aes_key = derive_key(shared_secret, iv)

    header = build_header(len(mlkem_ct), IV_LEN, TAG_LEN)

    cipher = Cipher(algorithms.AES(aes_key), modes.GCM(iv))
    encryptor = cipher.encryptor()
    # Preamble (header + ML-KEM CT + nonce) is authenticated as GCM AAD
    preamble = header + mlkem_ct + iv
    encryptor.authenticate_additional_data(preamble)

    total_size = os.path.getsize(plaintext_path)
    written = 0

    try:
        with open(output_path, 'wb') as fout:
            fout.write(preamble)

            with open(plaintext_path, 'rb') as fin:
                while True:
                    chunk = fin.read(CHUNK_SIZE)
                    if not chunk:
                        break
                    ct_chunk = encryptor.update(chunk)
                    fout.write(ct_chunk)
                    written += len(chunk)
                    print(f"  Progress: {written}/{total_size} bytes", end='\r')

            # finalize() computes the tag over AAD + all ciphertext
            encryptor.finalize()
            fout.write(encryptor.tag)
    except BaseException:
        # Never leave a partial .kyb behind on failure
        if os.path.exists(output_path):
            os.remove(output_path)
        raise

    print(f"\nSuccess: {plaintext_path} → {output_path}  ({total_size} bytes)")


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print(__doc__)
        sys.exit(1)

    infile = sys.argv[1]

    if not os.path.exists(infile):
        print(f"Error: File not found: {infile}")
        sys.exit(1)

    pubkey_path = "mlkem1024.pub"
    if not os.path.exists(pubkey_path):
        print(f"Error: No ML-KEM-1024 public key found. Run gen_keys.py first.")
        sys.exit(1)

    with open(pubkey_path, 'rb') as f:
        pub_key = f.read()

    outfile = infile + OUT_SUFFIX
    if os.path.exists(outfile):
        print(f"Error: Output already exists: {outfile}")
        sys.exit(1)

    encrypt_file(infile, outfile, pub_key)