137 lines
4.4 KiB
Python
137 lines
4.4 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Hybrid encrypt a file using AES-256-GCM + ML-KEM-1024.
|
|
Streaming handles files larger than RAM.
|
|
Usage:
|
|
python3 encrypt.py <filename>
|
|
|
|
Output: <filename>.kyb
|
|
Requires: mlkem1024.pub in the same directory.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import cbor2
|
|
import oqs
|
|
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
|
|
from cryptography.hazmat.primitives import hashes
|
|
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
|
|
|
|
MLKEM_KEM = "ML-KEM-1024"
|
|
SIGNATURE = b"KYB" # 3-byte file signature
|
|
VERSION = 1 # 1-byte format version
|
|
KDF_NAME = "HKDF-SHA256" # KDF advertised in the header
|
|
KDF_INFO = "mlkem-aes-key" # HKDF domain-separation label
|
|
OUT_SUFFIX = ".kyb"
|
|
IV_LEN = 12 # GCM standard nonce length (96 bits)
|
|
TAG_LEN = 16 # full GCM auth tag (128 bits)
|
|
CHUNK_SIZE = 64 * 1024 * 1024 # 64 MB
|
|
|
|
|
|
def build_header(mlkem_ct_len: int, iv_len: int, tag_len: int) -> bytes:
|
|
"""Build the CBOR header: signature + version + CBOR map of fields.
|
|
|
|
CBOR (RFC 8949) is compact, self-delimiting, and standardized in IETF
|
|
crypto formats (e.g. COSE), so the header stays forward-compatible.
|
|
"""
|
|
header_map = {
|
|
"KEM": "MLKEM1024",
|
|
"CIPHER": "AES256GCM",
|
|
"MLKEMCT": mlkem_ct_len,
|
|
"IV": iv_len,
|
|
"TAG": tag_len,
|
|
"KDF": "HKDF-SHA256",
|
|
"INFO": "mlkem-aes-key"
|
|
}
|
|
return SIGNATURE + bytes([VERSION]) + cbor2.dumps(header_map, canonical=True)
|
|
|
|
|
|
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 encrypt_file(plaintext_path: str, output_path: str, pub_key: bytes):
|
|
iv = os.urandom(IV_LEN)
|
|
|
|
# Each file gets a brand-new AES key:
|
|
# fresh shared secret (encap_secret) -> fresh derived key (derive_key).
|
|
# Since every file has its own unique key, even a repeated nonce between
|
|
# files would be safe. Never cache or reuse a shared secret across files -
|
|
# doing so would break that guarantee.
|
|
kem = oqs.KeyEncapsulation(MLKEM_KEM)
|
|
mlkem_ct, shared_secret = kem.encap_secret(pub_key)
|
|
kem.free()
|
|
|
|
aes_key = derive_key(shared_secret, iv)
|
|
|
|
header = build_header(len(mlkem_ct), IV_LEN, TAG_LEN)
|
|
|
|
cipher = Cipher(algorithms.AES(aes_key), modes.GCM(iv))
|
|
encryptor = cipher.encryptor()
|
|
# Preamble (header + ML-KEM CT + nonce) is authenticated as GCM AAD
|
|
preamble = header + mlkem_ct + iv
|
|
encryptor.authenticate_additional_data(preamble)
|
|
|
|
total_size = os.path.getsize(plaintext_path)
|
|
written = 0
|
|
|
|
try:
|
|
with open(output_path, 'wb') as fout:
|
|
fout.write(preamble)
|
|
|
|
with open(plaintext_path, 'rb') as fin:
|
|
while True:
|
|
chunk = fin.read(CHUNK_SIZE)
|
|
if not chunk:
|
|
break
|
|
ct_chunk = encryptor.update(chunk)
|
|
fout.write(ct_chunk)
|
|
written += len(chunk)
|
|
print(f" Progress: {written}/{total_size} bytes", end='\r')
|
|
|
|
# finalize() computes the tag over AAD + all ciphertext
|
|
encryptor.finalize()
|
|
fout.write(encryptor.tag)
|
|
except BaseException:
|
|
# Never leave a partial .kyb behind on failure
|
|
if os.path.exists(output_path):
|
|
os.remove(output_path)
|
|
raise
|
|
|
|
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 = "mlkem1024.pub"
|
|
if not os.path.exists(pubkey_path):
|
|
print(f"Error: No ML-KEM-1024 public key found. Run gen_keys.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)
|
|
|