Python simple KEM-TRANS example

demo_kem_trans.py

#!/usr/bin/env python3
from Crypto.Cipher import AES # https://pypi.org/project/pycryptodome/
from Crypto.Hash import SHAKE256

from types import SimpleNamespace


def _demo():
    # DEMO
    from pqc.kem import mceliece6960119f as kemalg # https://pypi.org/project/pypqc/
    from os import urandom

    pk, sk = kemalg.keypair()

    # kem-trans supports wrapping arbitrary sized data blobs
    data = urandom(69)
    # however, you should probably not use it
    # to wrap your *actual* message
    # because that will be inefficient (non streamable)
    # when using e.g. recommended SIV as wrapalg;
    # instead, just use kem-trans to wrap your actual
    # symmetric key which you use for hybrid encryption

    ek = kem_trans_wrap(kemalg, data, pk)

    result = kem_trans_unwrap(kemalg, ek, sk)

    if data == result:
        print("OK")
    else:
        raise ValueError("decryption error") # this should seriously NEVER occur


def siv_encrypt(k, m, *, aad=None, _alg=AES):
    # https://datatracker.ietf.org/doc/html/rfc5297#section-2.6
    cipher = _alg.new(k, _alg.MODE_SIV)
    if aad is not None:
        cipher.update(aad)
    ct, iv = cipher.encrypt_and_digest(m)
    assert len(iv) == _alg.block_size # SIV uses cipher's native block size, regardless of key size
    return iv + ct


def siv_decrypt(k, ct_siv, *, aad=None, _alg=AES):
    # https://datatracker.ietf.org/doc/html/rfc5297#section-2.7
    iv, ct = ct_siv[:_alg.block_size], ct_siv[_alg.block_size:]
    cipher = _alg.new(k, _alg.MODE_SIV)
    if aad is not None:
        cipher.update(aad)
    return cipher.decrypt_and_verify(ct, iv)


aes_siv = SimpleNamespace(encrypt=siv_encrypt, decrypt=siv_decrypt)


def shake256_256(m):
    h = SHAKE256.new()
    h.update(m)
    return h.read(32)


def kem_trans_wrap(kemalg, k, pk, *, wrapalg=aes_siv, kdf=shake256_256):
    # https://www.ietf.org/archive/id/draft-perret-prat-lamps-cms-pq-kem-00.html#name-senders-operations
    ct, ss = kemalg.encap(pk)
    kek = kdf(ss)
    wk = wrapalg.encrypt(kek, k)
    ek = wk + ct
    return ek


def kem_trans_unwrap(kemalg, ek, sk, *, wrapalg=aes_siv, kdf=shake256_256):
    # https://www.ietf.org/archive/id/draft-perret-prat-lamps-cms-pq-kem-00.html#name-recipients-operations
    wk, ct = ek[:-kemalg._lib.CRYPTO_CIPHERTEXTBYTES], ek[-kemalg._lib.CRYPTO_CIPHERTEXTBYTES:]
    try:
        ss = kemalg.decap(ct, sk)
    except TypeError as e: raise
    except Exception as e: raise ValueError("decryption error") from None
    kek = kdf(ss)
    try:
        k = wrapalg.decrypt(kek, wk)
    except TypeError as e: raise
    except Exception as e: raise ValueError("decryption error") from None
    return k


if __name__ == '__main__':
    _demo()