Generate encrypted and signed tokens for sharing DB primary keys

public_id.py

from hashlib import blake2s
from hmac import compare_digest
from typing import NamedTuple


def encode(key: bytes, value: int, expiry: int, key_id=0) -> str:
    """
    Encrypt and MAC a value + expiry + key_id.

    Params:
    - key: bytes used along with key_id to derive PRP and MAC keys
    - value: 64-bit unsigned int
    - key_id: 4-bit unsigned int
    - expiry: 32-bit unsigned int absolute timestamp (so max is 2106/02/07 17:28:15), minimum resolution is 16 seconds
      (rounds up). Set to 0 to never expire.

    Output is a str, that can be decrypted with the same key as above.
    """
    assert 0 <= value <= 0xFFFFFFFFFFFFFFFF, "value out of range"
    assert 0 <= expiry <= 0xFFFFFFFF, "expiry out of range"
    assert 0 <= key_id <= 0xF, "key_id out of range"
    """
    Binary layout:

    VVVVVVVV VVVVVVVV VVVVVVVV VVVVVVVV VVVVVVVV VVVVVVVV VVVVVVVV VVVVVVVV
    EEEEEEEE EEEEEEEE EEEEEEEE EEEEKKKK MMMMMMMM MMMMMMMM MMMMMMMM MMMMMMMM

    V value
    E expiry
    K key_id
    M mac
    """
    # Quantize expiry to 16 seconds (28-bits), rounding up
    expiry = (expiry + 15) >> 4

    # Derive keys
    key_id_b = key_id.to_bytes(1, "big")
    k_prp = _kdf(key, b"prp" + key_id_b)
    k_mac = _kdf(key, b"mac" + key_id_b)

    # Encode
    plaintext = (value << 28) | expiry

    # Encrypt
    prp = ((_feistel92_enc(plaintext, k_prp) << 4) | key_id).to_bytes(12, "big")

    # Sign
    mac = _mac(4, k_mac, prp)

    # Combine
    data = prp + mac

    return data.hex()


class DecodeResult(NamedTuple):
    value: int
    expiry: int


def decode_key_id(string: str) -> int:
    try:
        return int(string[23], 16)
    except Exception:
        return -1


def decode(key: bytes, string: str) -> DecodeResult:
    # Do not raise immediately, to reduce timing attacks
    error: Exception | None = None

    # Split prp and mac
    if len(string) != 32:
        error = ValueError(f"invalid length: expected 32 got {len(string)}")
    try:
        data = bytes.fromhex(string)
    except ValueError as e:
        error = e
        data = b""
    prp = data[:12]
    mac = data[12:]

    # Extract key_id
    key_id = decode_key_id(string)
    key_id_b = key_id.to_bytes(1, "big", signed=True)

    # Derive keys
    k_prp = _kdf(key, b"prp" + key_id_b)
    k_mac = _kdf(key, b"mac" + key_id_b)

    # Check signature
    mac_exp = _mac(4, k_mac, prp)
    if not compare_digest(mac_exp, mac) and not error:
        error = ValueError("invalid signature")

    # Decrypt
    plaintext = _feistel92_dec(int.from_bytes(prp, "big") >> 4, k_prp)

    # Decode
    expiry = (plaintext & 0xFFFFFFF) << 4
    value = plaintext >> 28

    # Return
    if error:
        raise error
    return DecodeResult(value, expiry)


_U46 = (1 << 46) - 1
_U92 = (1 << 92) - 1
_ROUNDS = 8
_ROUNDS_LABEL = b"feistel" + _ROUNDS.to_bytes(1, "big")


def _kdf(seed: bytes, label: bytes) -> bytes:
    return blake2s(label, key=seed, digest_size=16).digest()


def _mac(digest_size: int, key: bytes, *data: bytes) -> bytes:
    h = blake2s(digest_size=digest_size, key=key)
    for d in data:
        h.update(d)
    return h.digest()


def _feistel92_round(r46: int, key: bytes, rnd: int) -> int:
    return int.from_bytes(
        _mac(6, key, _ROUNDS_LABEL, ((r46 << 8) | rnd).to_bytes(7, "big")), "big"
    )


def _feistel92_enc(x: int, key: bytes) -> int:
    L = (x >> 46) & _U46
    R = x & _U46
    for i in range(1, _ROUNDS + 1):
        L, R = R, (L ^ _feistel92_round(R, key, i)) & _U46
    return ((L << 46) | R) & _U92


def _feistel92_dec(y: int, key: bytes) -> int:
    L = (y >> 46) & _U46
    R = y & _U46
    for i in range(_ROUNDS, 0, -1):
        L, R = (R ^ _feistel92_round(L, key, i)) & _U46, L
    return ((L << 46) | R) & _U92