Generate encrypted and signed tokens for sharing DB primary keys

public_id.py

"""
Encoding scheme for obfuscating 64-bit integer IDs into 32-character signed hexadecimal tokens.

Includes support for optional expiry timestamps, and is designed to be fast and secure against guessing or tampering.

The same (key, value, expiry) always produces the same token, and small changes to any of those inputs produce
completely different tokens.

Uses HMAC with SHA-256 so it can be easily implemented in other languages.
"""

import hmac
import struct
import time
from datetime import datetime
from functools import cache


def encode(key: bytes, value: int, expiry: datetime | None = None) -> str:
    if not (0 <= value <= 0xFFFFFFFFFFFFFFFF):
        raise ValueError(f"out of range: {value}")
    expiry_ts = int(expiry.timestamp()) if expiry else 0
    if not (0 <= expiry_ts <= 0xFFFFFFFF):
        raise ValueError(f"expiry timestamp out of range: {expiry_ts}")
    k_mac, k_pad = _kdf(key)
    c = struct.pack(">QI", value, expiry_ts)
    tag = hmac.new(k_mac, c, "sha256").digest()[:4]
    pad = hmac.new(k_pad, tag, "sha256").digest()[:12]
    msg = (int.from_bytes(c, "big") ^ int.from_bytes(pad, "big")).to_bytes(12, "big")
    return (msg + tag).hex()


def decode(key: bytes, token: str) -> int:
    try:
        raw = bytes.fromhex(token)
    except ValueError as e:
        raise ValueError(f"token is not valid hex: {token!r}") from e
    if len(raw) != 16:
        raise ValueError(f"invalid token length: {len(raw)} != 16")
    k_mac, k_pad = _kdf(key)
    msg = raw[:12]
    tag = raw[12:16]
    pad = hmac.new(k_pad, tag, "sha256").digest()[:12]
    c = (int.from_bytes(msg, "big") ^ int.from_bytes(pad, "big")).to_bytes(12, "big")
    expected_tag = hmac.new(k_mac, c, "sha256").digest()[:4]
    if not hmac.compare_digest(tag, expected_tag):
        raise ValueError("invalid token")
    value, expiry_ts = struct.unpack(">QI", c)
    if expiry_ts != 0 and expiry_ts < int(time.time()):
        raise ValueError("token has expired")
    return value


@cache
def _kdf(key: bytes) -> tuple[bytes, bytes]:
    k_mac = hmac.new(key, b"MAC\x01", "sha256").digest()
    k_pad = hmac.new(key, b"PAD\x01", "sha256").digest()
    return k_mac, k_pad