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

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

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

KYBER_KEM    = "Kyber1024"
SIGNATURE    = "KYB"
VERSION      = "1"
OUT_SUFFIX   = ".kyb"
IV_LEN       = 16
TAG_LEN      = 32
CHUNK_SIZE   = 64 * 1024 * 1024   # 64 MB


def build_header(kyber_ct_len: int, iv_len: int, tag_len: int) -> bytes:
    """Build the self-describing ASCII header line (ends with a newline).

    Human-readable and labelled, so a decrypt script (or a person) can learn
    every format parameter straight from the file with no hardcoded sizes.
    """
    fields = [
        SIGNATURE,
        f"VER={VERSION}",
        "KEM=MLKEM1024",
        "CIPHER=AES256CTR",
        "MAC=HMACSHA256",
        f"KYBERCT={kyber_ct_len}",
        f"IV={iv_len}",
        f"TAG={tag_len}",
    ]
    return "|".join(fields).encode("ascii") + b"\n"


def derive_keys(shared_secret: bytes, salt: bytes):
    """Derive independent AES and HMAC keys from the Kyber shared secret.

    The IV is used as the HKDF salt, binding the keys to this specific
    encryption (defense-in-depth: key/nonce mismatches fail loudly).
    """
    aes_key = HKDF(
        algorithm=hashes.SHA256(), length=32, salt=salt,
        info=b"kyber-aes-key"
    ).derive(shared_secret)
    hmac_key = HKDF(
        algorithm=hashes.SHA256(), length=32, salt=salt,
        info=b"kyber-hmac-key"
    ).derive(shared_secret)
    return aes_key, hmac_key


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

    # Fresh encapsulation per file -> fresh shared secret -> fresh AES/HMAC keys.
    # Never reuse or cache a shared secret across files: the per-file keys are
    # exactly what make accidental IV reuse harmless.
    kem = oqs.KeyEncapsulation(KYBER_KEM)
    kyber_ct, shared_secret = kem.encap_secret(pub_key)
    kem.free()

    aes_key, hmac_key = derive_keys(shared_secret, iv)

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

    cipher = Cipher(algorithms.AES(aes_key), modes.CTR(iv))
    encryptor = cipher.encryptor()
    hmac = HMAC(hmac_key, hashes.SHA256())

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

    with open(output_path, 'wb') as fout:
        # Preamble (all authenticated): ASCII header + Kyber CT + IV
        preamble = header + kyber_ct + iv
        fout.write(preamble)
        hmac.update(preamble)

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

            final_ct = encryptor.finalize()
            hmac.update(final_ct)   # cover every ciphertext byte written (no-op for CTR)
            fout.write(final_ct)
            fout.write(hmac.finalize())

    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 = "kyber.pub"
    if not os.path.exists(pubkey_path):
        print(f"Error: No Kyber public key found. Run gen_kyber.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)