#!/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>.enc
Requires: kyber.pub in the same directory.
"""

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

KYBER_KEM    = "Kyber1024"
IV_LEN       = 16
CHUNK_SIZE   = 64 * 1024 * 1024   # 64 MB
HMAC_KEY_LEN = 32


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

    kem = oqs.KeyEncapsulation(KYBER_KEM)
    kyber_ct, aes_key = kem.encap_secret(pub_key)
    kem.free()

    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:
        # Header
        fout.write(struct.pack('!I', len(kyber_ct)))
        fout.write(kyber_ct)
        fout.write(iv)
        fout.write(hmac_key)

        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')

            fout.write(encryptor.finalize())
            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 + ".enc"
    encrypt_file(infile, outfile, pub_key)