kyb/decrypt.py

#!/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.kyb>

Output: original filename (without .kyb)
Requires: kyber.sec in the same directory.
"""

import os
import sys
import tempfile
import hmac as hmac_mod
import oqs
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.hmac import HMAC
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

SIGNATURE      = "KYB"
SUPPORTED_VER  = "1"
INPUT_SUFFIX   = ".kyb"
CHUNK_SIZE     = 64 * 1024 * 1024
MAX_HEADER_LEN = 1024    # sane upper bound for the header line


def derive_keys(shared_secret: bytes, salt: bytes):
    """Derive independent AES and HMAC keys from the Kyber shared secret.

    The IV is used as the HKDF salt, binding the keys to this specific
    encryption (defense-in-depth: key/nonce mismatches fail loudly).
    """
    aes_key = HKDF(
        algorithm=hashes.SHA256(), length=32, salt=salt,
        info=b"kyber-aes-key"
    ).derive(shared_secret)
    hmac_key = HKDF(
        algorithm=hashes.SHA256(), length=32, salt=salt,
        info=b"kyber-hmac-key"
    ).derive(shared_secret)
    return aes_key, hmac_key


def parse_header(header_bytes: bytes):
    """Parse the ASCII header line into a dict of fields."""
    line = header_bytes.rstrip(b"\n").decode("ascii")
    parts = line.split("|")
    if parts[0] != SIGNATURE:
        raise ValueError("Error: Not a kyber (.kyb) file (bad signature)")
    fields = {}
    for part in parts[1:]:
        if "=" in part:
            key, _, value = part.partition("=")
            fields[key] = value
    return fields


def decrypt_file(input_path: str, output_path: str, secret_key: bytes):
    with open(input_path, 'rb') as fin:
        # 1. Parse the self-describing header line
        header_bytes = fin.readline(MAX_HEADER_LEN)
        if not header_bytes.endswith(b"\n"):
            raise ValueError("Error: File truncated or corrupt header")
        fields = parse_header(header_bytes)

        # 2. Version gate: refuse files this tool cannot parse
        version = fields.get("VER")
        if version != SUPPORTED_VER:
            raise ValueError(f"Error: Unsupported format version: {version}")

        # 3. 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") != "AES256CTR":
            raise ValueError(f"Error: Unsupported cipher: {fields.get('CIPHER')}")
        if fields.get("MAC") != "HMACSHA256":
            raise ValueError(f"Error: Unsupported MAC: {fields.get('MAC')}")

        try:
            kyber_ct_len = int(fields["KYBERCT"])
            iv_len = int(fields["IV"])
            tag_len = int(fields["TAG"])
        except (KeyError, ValueError):
            raise ValueError("Error: Corrupt header (missing or invalid length fields)")

        # Cross-checks: ML-KEM-1024 CT is always 1568 B; HMAC-SHA256 tag always 32 B
        if kyber_ct_len != 1568:
            raise ValueError("Error: Header inconsistent: MLKEM1024 implies a 1568-byte ciphertext")
        if tag_len != 32:
            raise ValueError("Error: Header inconsistent: HMACSHA256 implies a 32-byte tag")

        # 4. Read the Kyber ciphertext + IV
        kyber_ct = fin.read(kyber_ct_len)
        if len(kyber_ct) != kyber_ct_len:
            raise ValueError("Error: File truncated: incomplete Kyber ciphertext")
        iv = fin.read(iv_len)
        if len(iv) != iv_len:
            raise ValueError("Error: File truncated: incomplete header")

        # 5. Recover the shared secret, then derive AES + HMAC keys
        kem = oqs.KeyEncapsulation("Kyber1024", secret_key)
        try:
            shared_secret = kem.decap_secret(kyber_ct)
        except Exception:
            raise ValueError("Error: Failed to decapsulate key - wrong secret key or corrupt header")
        finally:
            kem.free()
        aes_key, hmac_key = derive_keys(shared_secret, iv)

        # 6. Stream-decrypt, authenticating the whole preamble + ciphertext
        cipher = Cipher(algorithms.AES(aes_key), modes.CTR(iv))
        decryptor = cipher.decryptor()
        hmac = HMAC(hmac_key, hashes.SHA256())
        hmac.update(header_bytes)
        hmac.update(kyber_ct)
        hmac.update(iv)

        # File size minus the MAC tag at the end
        fin.seek(0, os.SEEK_END)
        total_size = fin.tell()
        if total_size < len(header_bytes) + kyber_ct_len + iv_len + tag_len:
            raise ValueError("Error: File too small to be a valid .kyb file")

        data_end = total_size - tag_len
        body_start = len(header_bytes) + kyber_ct_len + iv_len
        fin.seek(body_start)
        total_read = data_end - body_start

        # Write to a temp file; only rename into place after the HMAC 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)
                    hmac.update(chunk)
                    fout.write(decryptor.update(chunk))
                    read += len(chunk)
                    print(f"  Progress: {read}/{total_read} bytes", end='\r')

                # Verify integrity (constant-time comparison)
                stored_hmac = fin.read(tag_len)
                computed_hmac = hmac.finalize()
                if not hmac_mod.compare_digest(stored_hmac, computed_hmac):
                    print()
                    raise ValueError("Error: Integrity check failed. File corrupted or tampered")

                fout.write(decryptor.finalize())
                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 = "kyber.sec"
    if not os.path.exists(seckey_path):
        print("Error: No Kyber 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)