nordisk/gunnar

crates/gunnar-crypt/tests/vectors/check-v2.py

#!/usr/bin/env python3
"""The envelope v2 vectors (cases-v2.txt), from an INDEPENDENT implementation.

A second implementation of the v2 format from its description in
`.nornir/encrypted-repos.md` §4.1, on Python's `cryptography` (AESSIV =
RFC 5297, HKDF, HMAC — no gunnar code). `--write` makes cases-v2.txt, which
`tests/chunked.rs` then checks the Rust build against, native and wasm32;
without it, it re-checks the file (seal every case, open every spelled-out
envelope). TEST KEYS ONLY (vectors/keyring.txt).

    python3 crates/gunnar-crypt/tests/vectors/check-v2.py [--write]
"""
import base64
import hashlib
import hmac
import os
import struct
import sys

from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers.aead import AESSIV
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

HERE = os.path.dirname(os.path.abspath(__file__))
MAGIC = b"\0GNCRYPT"
NONCE_INFO = b"gunnar-crypt envelope v2 file nonce"
CHUNK = 1 << 20


def ring():
    keys = {}
    for line in open(os.path.join(HERE, "keyring.txt")):
        f = line.split()
        if len(f) == 3 and f[0] == "epoch":
            keys[int(f[1])] = base64.b64decode(f[2])
    return keys


def generated(n):
    return bytes(
        ((((i * 2654435761) & 0xFFFFFFFF) >> 24) ^ (i >> 20)) & 0xFF for i in range(n)
    )


def nonce(key, plain):
    nk = HKDF(algorithm=hashes.SHA256(), length=32, salt=None, info=NONCE_INFO).derive(key)
    return hmac.new(nk, plain, hashlib.sha256).digest()


def seal(epoch, key, chunk, plain):
    header = MAGIC + bytes([2]) + struct.pack(">II", epoch, chunk) + nonce(key, plain)
    siv = AESSIV(key)
    pieces = [plain[i : i + chunk] for i in range(0, len(plain), chunk)] or [b""]
    out = [header]
    for i, p in enumerate(pieces):
        last = 1 if i == len(pieces) - 1 else 0
        ad = header + struct.pack(">Q", i) + bytes([last])
        out.append(siv.encrypt(p, [ad]))  # tag ‖ ciphertext
    return b"".join(out)


def open_v2(keys, env):
    assert env[:8] == MAGIC and env[8] == 2
    epoch, chunk = struct.unpack(">II", env[9:17])
    header, body = env[:49], env[49:]
    siv = AESSIV(keys[epoch])
    blocks = [body[i : i + chunk + 16] for i in range(0, len(body), chunk + 16)]
    out = []
    for i, b in enumerate(blocks):
        last = 1 if i == len(blocks) - 1 else 0
        out.append(siv.decrypt(b, [header + struct.pack(">Q", i) + bytes([last])]))
    return b"".join(out)


SPELLED = [
    # name, epoch, chunk, plaintext
    ("empty", 1, CHUNK, b""),
    ("hello", 1, CHUNK, b"hello, world\n"),
    ("binary", 1, CHUNK, bytes([0, 255, 1, 254, 0, 0, 7])),
    ("utf8", 2, CHUNK, "räksmörgås ✓\n".encode()),
    ("long", 2, CHUNK, b"x" * 4099),
    ("blocks-exact", 1, 16, generated(64)),
    ("blocks-ragged", 2, 16, generated(69)),
    ("one-byte-chunks", 1, 1, generated(5)),
]
GENERATED = [
    # name, epoch, plaintext length (chunk = 1 MiB)
    ("one-mib", 1, CHUNK),
    ("one-mib-and-a-byte", 2, CHUNK + 1),
    ("three-and-a-half-mib", 2, 3 * CHUNK + CHUNK // 2),
]


def write(keys):
    lines = [
        "# kind name epoch chunk plaintext envelope — TEST KEYS ONLY (vectors/keyring.txt)",
        "# written by check-v2.py --write (Python cryptography, not gunnar's code)",
        "# hex: plaintext base64 (- = empty), envelope hex",
        "# gen: plaintext = generated(len) in tests/chunked.rs, envelope = length:sha256",
    ]
    for name, epoch, chunk, plain in SPELLED:
        env = seal(epoch, keys[epoch], chunk, plain)
        p = base64.b64encode(plain).decode() if plain else "-"
        lines.append(f"hex {name} {epoch} {chunk} {p} {env.hex()}")
    for name, epoch, n in GENERATED:
        env = seal(epoch, keys[epoch], CHUNK, generated(n))
        lines.append(f"gen {name} {epoch} {CHUNK} {n} {len(env)}:{hashlib.sha256(env).hexdigest()}")
    with open(os.path.join(HERE, "cases-v2.txt"), "w") as f:
        f.write("\n".join(lines) + "\n")


def main():
    keys = ring()
    if "--write" in sys.argv[1:]:
        write(keys)
    n = 0
    for line in open(os.path.join(HERE, "cases-v2.txt")):
        if not line.strip() or line.startswith("#"):
            continue
        f = line.split()
        kind, name, epoch, chunk = f[0], f[1], int(f[2]), int(f[3])
        if kind == "hex":
            plain = b"" if f[4] == "-" else base64.b64decode(f[4])
            want = bytes.fromhex(f[5])
            got = seal(epoch, keys[epoch], chunk, plain)
            assert got == want, f"{name}: the Rust bytes differ from the Python seal"
            assert open_v2(keys, want) == plain, f"{name}: does not open in Python"
        else:
            plain = generated(int(f[4]))
            length, digest = f[5].split(":")
            got = seal(epoch, keys[epoch], chunk, plain)
            assert len(got) == int(length), f"{name}: length"
            assert hashlib.sha256(got).hexdigest() == digest, f"{name}: digest"
            assert open_v2(keys, got) == plain
        n += 1
    print(f"PASS: {n} v2 vectors agree with an independent implementation (Python cryptography AESSIV)")
    return 0


if __name__ == "__main__":
    sys.exit(main())