198 lines
7.7 KiB
Python
198 lines
7.7 KiB
Python
#!/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 <filename.kyb>
|
|
|
|
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)
|
|
|