#!/usr/bin/env python3 """ Hybrid decrypt a file encrypted with AES-256-GCM + ML-KEM-1024. Streaming handles files larger than RAM. Usage: python3 decrypt.py Output: original filename (without .kyb) Requires: mlkem1024.sec in the same directory. """ import os import sys import tempfile from io import BytesIO import cbor2 import oqs from cryptography.exceptions import InvalidTag from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes from cryptography.hazmat.primitives import hashes from cryptography.hazmat.primitives.kdf.hkdf import HKDF SIGNATURE = b"KYB" SUPPORTED_VER = 1 INPUT_SUFFIX = ".kyb" CHUNK_SIZE = 64 * 1024 * 1024 MAX_HEADER_LEN = 1024 # sane upper bound for the CBOR header map HEADER_PREFIX_LEN = 4 # signature (3) + version (1) KDF_NAME = "HKDF-SHA256" # KDF advertised in the header KDF_INFO = "mlkem-aes-key" # HKDF domain-separation label 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 decrypt_file(input_path: str, output_path: str, secret_key: bytes): with open(input_path, 'rb') as fin: # 1. Read the fixed prefix: signature (3) + version (1) prefix = fin.read(HEADER_PREFIX_LEN) if len(prefix) != HEADER_PREFIX_LEN: raise ValueError("Error: File truncated or corrupt header") if prefix[:3] != SIGNATURE: raise ValueError("Error: Not a .kyb file (bad signature)") # 2. Version gate: refuse files this tool cannot parse if prefix[3] != SUPPORTED_VER: raise ValueError(f"Error: Unsupported format version: {prefix[3]}") # 3. Read the CBOR header map (bounded read, then locate its end) chunk = fin.read(MAX_HEADER_LEN) if not chunk: raise ValueError("Error: File truncated or corrupt header") decoder = cbor2.CBORDecoder(BytesIO(chunk)) try: fields = decoder.decode() cbor_len = decoder.fp.tell() except EOFError: raise ValueError("Error: File truncated or corrupt header") except Exception: raise ValueError("Error: Corrupt header (invalid CBOR)") if not isinstance(fields, dict): raise ValueError("Error: Corrupt header (CBOR map expected)") header_bytes = prefix + chunk[:cbor_len] # Position exactly after the CBOR map so the ML-KEM CT reads correctly fin.seek(HEADER_PREFIX_LEN + cbor_len) # 4. Validate the advertised algorithms and derive the lengths kem_name = fields.get("KEM") if kem_name not in ("MLKEM1024",): raise ValueError(f"Error: Unsupported KEM: {kem_name}") if fields.get("CIPHER") != "AES256GCM": raise ValueError(f"Error: Unsupported cipher: {fields.get('CIPHER')}") if fields.get("KDF") != KDF_NAME: raise ValueError(f"Error: Unsupported KDF: {fields.get('KDF')}") if fields.get("INFO") != KDF_INFO: raise ValueError(f"Error: Unsupported KDF info: {fields.get('INFO')}") try: mlkem_ct_len = fields["MLKEMCT"] iv_len = fields["IV"] tag_len = fields["TAG"] except (KeyError, TypeError): raise ValueError("Error: Corrupt header (missing or invalid length fields)") # Cross-checks: ML-KEM-1024 CT is always 1568 B; GCM nonce 12 B; tag 16 B if mlkem_ct_len != 1568: raise ValueError("Error: Header inconsistent: MLKEM1024 implies a 1568-byte ciphertext") if iv_len != 12: raise ValueError("Error: Header inconsistent: AES256GCM implies a 12-byte nonce") if tag_len != 16: raise ValueError("Error: Header inconsistent: AES256GCM implies a 16-byte tag") # 5. Read the ML-KEM ciphertext + IV mlkem_ct = fin.read(mlkem_ct_len) if len(mlkem_ct) != mlkem_ct_len: raise ValueError("Error: File truncated: incomplete ML-KEM ciphertext") iv = fin.read(iv_len) if len(iv) != iv_len: raise ValueError("Error: File truncated: incomplete header") # 6. Recover the shared secret, then derive the AES key kem = oqs.KeyEncapsulation("ML-KEM-1024", secret_key) try: shared_secret = kem.decap_secret(mlkem_ct) except Exception: raise ValueError("Error: Failed to decapsulate key - wrong secret key or corrupt header") finally: kem.free() aes_key = derive_key(shared_secret, iv) # File size minus the GCM tag at the end fin.seek(0, os.SEEK_END) total_size = fin.tell() if total_size < len(header_bytes) + mlkem_ct_len + iv_len + tag_len: raise ValueError("Error: File too small to be a valid .kyb file") # The trailing GCM tag must be supplied up front (pyca/cryptography) fin.seek(total_size - tag_len) gcm_tag = fin.read(tag_len) data_end = total_size - tag_len body_start = len(header_bytes) + mlkem_ct_len + iv_len fin.seek(body_start) total_read = data_end - body_start # 7. Stream-decrypt, authenticating the whole preamble + ciphertext cipher = Cipher(algorithms.AES(aes_key), modes.GCM(iv, gcm_tag)) decryptor = cipher.decryptor() decryptor.authenticate_additional_data(header_bytes + mlkem_ct + iv) # Write to a temp file; only rename into place after the tag verifies, # so corrupted plaintext never appears at the destination path. fd, temp_path = tempfile.mkstemp(dir=os.path.dirname(output_path) or ".", prefix=".dec-", suffix=".tmp") try: with os.fdopen(fd, 'wb') as fout: read = 0 while fin.tell() < data_end: chunk_size = min(CHUNK_SIZE, data_end - fin.tell()) chunk = fin.read(chunk_size) fout.write(decryptor.update(chunk)) read += len(chunk) print(f" Progress: {read}/{total_read} bytes", end='\r') # Verify integrity: finalize() checks the GCM tag try: decryptor.finalize() except InvalidTag: print() raise ValueError("Error: Integrity check failed. File corrupted or tampered") fout.flush() os.fsync(fout.fileno()) # Atomic move: only a fully verified file lands at output_path os.replace(temp_path, output_path) except BaseException: if os.path.exists(temp_path): os.remove(temp_path) raise 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(INPUT_SUFFIX): print(f"Warning: File doesn't end with {INPUT_SUFFIX}. Decrypting anyway") seckey_path = "mlkem1024.sec" if not os.path.exists(seckey_path): print("Error: No ML-KEM-1024 secret key found.") sys.exit(1) with open(seckey_path, 'rb') as f: secret_key = f.read() outfile = infile[:-len(INPUT_SUFFIX)] if infile.endswith(INPUT_SUFFIX) else infile + ".dec" decrypt_file(infile, outfile, secret_key)