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

Output: original filename (without .enc)
Requires: kyber.sec 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
HMAC_KEY_LEN  = 32
HMAC_TAG_LEN  = 32


def decrypt_file(input_path: str, output_path: str, secret_key: bytes):
    with open(input_path, 'rb') as fin:
        # Read header
        kyber_ct_len = struct.unpack('!I', fin.read(4))[0]
        kyber_ct     = fin.read(kyber_ct_len)
        iv           = fin.read(IV_LEN)
        hmac_key     = fin.read(HMAC_KEY_LEN)

        # Recover AES key
        kem = oqs.KeyEncapsulation(KYBER_KEM, secret_key)
        aes_key = kem.decap_secret(kyber_ct)
        kem.free()

        # Decrypt stream
        cipher = Cipher(algorithms.AES(aes_key), modes.CTR(iv))
        decryptor = cipher.decryptor()
        hmac = HMAC(hmac_key, hashes.SHA256())

        # File size minus HMAC tag at the end
        fin.seek(0, os.SEEK_END)
        total_size = fin.tell()
        data_end = total_size - HMAC_TAG_LEN
        fin.seek(4 + kyber_ct_len + IV_LEN + HMAC_KEY_LEN)

        total_read = data_end - fin.tell()

        with open(output_path, 'wb') as fout:
            read = 0
            while fin.tell() < data_end:
                chunk_size = min(CHUNK_SIZE, data_end - fin.tell())
                chunk = fin.read(chunk_size)
                hmac.update(chunk)
                fout.write(decryptor.update(chunk))
                read += len(chunk)
                print(f"  Progress: {read}/{total_read} bytes", end='\r')

            # Verify integrity
            stored_hmac = fin.read(HMAC_TAG_LEN)
            computed_hmac = hmac.finalize()
            if stored_hmac != computed_hmac:
                os.remove(output_path)
                print()
                raise ValueError("Error: Integrity check failed. File corrupted or tampered")

            fout.write(decryptor.finalize())

    print(f"\nSuccess: {input_path} → {output_path}  ({total_read} 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)

    if not infile.endswith(".enc"):
        print(f"Warning:  File doesn't end with .enc. Decrypting anyway")

    seckey_path = "kyber.sec"
    if not os.path.exists(seckey_path):
        print(f"Error: No Kyber secret key found.")
        sys.exit(1)

    with open(seckey_path, 'rb') as f:
        secret_key = f.read()

    outfile = infile[:-4] if infile.endswith(".enc") else infile + ".dec"
    decrypt_file(infile, outfile, secret_key)