380 lines
16 KiB
Python
380 lines
16 KiB
Python
"""Hybrid ML-KEM-768 + X25519 mutually authenticated key exchange on a
|
|
symmetric shared key, chab3-boxed throughout. Pure cryptography: a
|
|
Registration or Resumption object holds one exchange's per-connection
|
|
state, disposed by the caller at completion; sansio.py owns the state
|
|
machine, key store and policies. See docs/protocol.md.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hmac
|
|
import struct
|
|
|
|
from cryptography.hazmat.primitives.asymmetric.mlkem import (
|
|
MLKEM768PrivateKey,
|
|
MLKEM768PublicKey,
|
|
)
|
|
from cryptography.hazmat.primitives.asymmetric.x25519 import (
|
|
X25519PrivateKey,
|
|
X25519PublicKey,
|
|
)
|
|
|
|
from .chab3 import decrypt, encrypt, random
|
|
from .exceptions import _Burn
|
|
from .kdf import kdf
|
|
|
|
# -- Fixed-size record layouts (struct handles packing + length checks) -----
|
|
# Hybrid scheme = ML-KEM-768 (pk 1184 / sk-seed 64 / ct 1088) || X25519 (32)
|
|
# Handshake message bodies; box(N) is a chab3 box of N plaintext bytes (N+28)
|
|
REG_MSG1 = struct.Struct("1184s32s32s") # e_alice_m | e_alice_x | commitment
|
|
REG_MSG2 = struct.Struct("32s1088s1244s") # e_bob | ee_ct | encrypt_ee(inner)
|
|
REG_MSG2_INNER = struct.Struct("1184s32s") # s_bob_m | s_bob_x
|
|
REG_MSG3 = struct.Struct("1088s1276s") # es_ct | encrypt_es(inner)
|
|
REG_MSG3_INNER = struct.Struct("1184s32s32s") # s_alice_m | s_alice_x | reveal
|
|
REG_MSG4 = struct.Struct("1088s32s") # se_ct | tb
|
|
# REG_MSG5 is plain tag, no Struct needed
|
|
|
|
RES_MSG1 = struct.Struct("1184s32s1088s") # e_alice_m | e_alice_x | sb_ct
|
|
RES_MSG2 = struct.Struct("32s1088s1088s32s") # e_bob | sa_ct | ee_ct | tb
|
|
|
|
|
|
def _mlkem_private_key(idroot: bytes) -> MLKEM768PrivateKey:
|
|
return MLKEM768PrivateKey.from_seed_bytes(kdf("idroot/mlkem", idroot, length=64))
|
|
|
|
|
|
def _x25519_private_key(idroot: bytes) -> X25519PrivateKey:
|
|
return X25519PrivateKey.from_private_bytes(kdf("idroot/x25519", idroot))
|
|
|
|
|
|
def _public_from_idroot(idroot: bytes) -> tuple[bytes, bytes]:
|
|
# (mlkem public, x25519 public) derived from the idroot
|
|
return (
|
|
_mlkem_private_key(idroot).public_key().public_bytes_raw(),
|
|
_x25519_private_key(idroot).public_key().public_bytes_raw(),
|
|
)
|
|
|
|
|
|
def _fresh_ephemeral() -> tuple[MLKEM768PrivateKey, X25519PrivateKey, bytes, bytes]:
|
|
# (mlkem key, x25519 key, mlkem public, x25519 public), fresh per connection
|
|
m_sk = MLKEM768PrivateKey.generate()
|
|
x_sk = X25519PrivateKey.generate()
|
|
return (
|
|
m_sk,
|
|
x_sk,
|
|
m_sk.public_key().public_bytes_raw(),
|
|
x_sk.public_key().public_bytes_raw(),
|
|
)
|
|
|
|
|
|
def _traffic_keys(key: bytes, initiator: bool) -> tuple[bytes, bytes]:
|
|
# -> (send, recv); same input format in both flows -> same labels
|
|
a2b = kdf("a2b", key)
|
|
b2a = kdf("b2a", key)
|
|
return (a2b, b2a) if initiator else (b2a, a2b)
|
|
|
|
|
|
def _check(ok: bool, why: str) -> None:
|
|
if not ok:
|
|
raise _Burn(why)
|
|
|
|
|
|
def open_hello(msg1: bytes, *, shared: bytes, ad: bytes, flow: str) -> bytes:
|
|
# Open a msg1 box under a candidate credential's shared key (the
|
|
# responder's caller scan); ValueError = not this credential
|
|
return decrypt(msg1, key=kdf(flow + "/msg1", shared), ad=ad)
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# AD-chained boxes: each box's associated data is the previous box's tag
|
|
|
|
|
|
class _Chain:
|
|
def __init__(self, ad: bytes):
|
|
self._ad = ad
|
|
|
|
def encrypt(self, body: bytes, *, key: bytes) -> bytes:
|
|
msg = encrypt(body, key=key, ad=self._ad)
|
|
self._ad = msg[-16:]
|
|
return bytes(msg)
|
|
|
|
def decrypt(self, msg: bytes, *, key: bytes) -> bytes:
|
|
plain = decrypt(msg, key=key, ad=self._ad)
|
|
self._ad = msg[-16:]
|
|
return plain
|
|
|
|
def advance(self, msg: bytes) -> None:
|
|
# A box decrypted outside the chain (the responder's caller scan)
|
|
self._ad = msg[-16:]
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# The two exchanges. Stage names carry the acting party:
|
|
# Alice drives with startA/respondA (+ confirmA when registering),
|
|
# Bob with respondB (+ confirmB/finishB when registering, completeB
|
|
# when resuming)
|
|
|
|
|
|
class Registration:
|
|
"""Registration exchange (first connection): every leg is a bare
|
|
ML-KEM ciphertext plus an X25519 DH against the shipped connection
|
|
keys (e_alice in msg1, e_bob in msg2), so related secrets provably
|
|
come from the same party. finish() -> (send, recv, rotated shared,
|
|
sas, peer pub)."""
|
|
|
|
def __init__(self, shared: bytes, idroot: bytes, ad: bytes):
|
|
self._shared = shared
|
|
self._idroot = idroot
|
|
# Static key objects, derived once from the persisted idroot
|
|
self._mlkem = _mlkem_private_key(idroot)
|
|
self._x25519 = _x25519_private_key(idroot)
|
|
self._static_m, self._static_x = _public_from_idroot(idroot)
|
|
self._chain = _Chain(ad)
|
|
|
|
def _root(self) -> bytes:
|
|
return kdf("reg/root", self._shared, self._ee, self._es, self._se)
|
|
|
|
# -- alice ----------------------------------------------------------------
|
|
|
|
def startA(self) -> bytes:
|
|
# msg1: our connection keys, with the reveal committed so it
|
|
# cannot be predicted or steered
|
|
self._eph_m, self._ephemeral, e_alice_m, e_alice_x = _fresh_ephemeral()
|
|
self._reveal = random(32)
|
|
commitment = kdf("commit", e_alice_m, e_alice_x, self._reveal)
|
|
body = REG_MSG1.pack(e_alice_m, e_alice_x, commitment)
|
|
return self._chain.encrypt(body, key=kdf("reg/msg1", self._shared))
|
|
|
|
def respondA(self, msg2: bytes) -> bytes:
|
|
# Open Bob's reply -- an outer AEAD failure burns nothing, it
|
|
# proves no key knowledge -- then challenge his static
|
|
plain = self._chain.decrypt(msg2, key=kdf("reg/msg2", self._shared))
|
|
try:
|
|
e_bob, ee_ct, inner = REG_MSG2.unpack(plain)
|
|
# ee is forward secrecy, made to our connection key only
|
|
ee = self._eph_m.decapsulate(ee_ct)
|
|
ee += self._ephemeral.exchange(X25519PublicKey.from_public_bytes(e_bob))
|
|
# Solving ee lets us read Bob's static from the inner box
|
|
s_bob_m, s_bob_x = REG_MSG2_INNER.unpack(
|
|
decrypt(inner, key=kdf("reg/msg2/inner", ee))
|
|
)
|
|
# We challenge his static: only the real Bob can solve this
|
|
es, es_ct = MLKEM768PublicKey.from_public_bytes(s_bob_m).encapsulate()
|
|
es += self._ephemeral.exchange(X25519PublicKey.from_public_bytes(s_bob_x))
|
|
except (struct.error, ValueError) as e:
|
|
raise _Burn(f"bad msg2 contents ({e})") from None
|
|
self._peer_pub = (s_bob_m, s_bob_x)
|
|
self._e_bob = e_bob
|
|
self._ee = ee
|
|
self._es = es
|
|
return self._chain.encrypt(
|
|
es_ct
|
|
+ encrypt(
|
|
REG_MSG3_INNER.pack(self._static_m, self._static_x, self._reveal),
|
|
key=kdf("reg/msg3/inner", es),
|
|
),
|
|
key=kdf("reg/msg3", self._shared, ee),
|
|
)
|
|
|
|
def confirmA(self, msg4: bytes) -> tuple[bytes, bytes]:
|
|
# Bob's challenge to our static plus his root confirmation. The
|
|
# caller gates the answer (msg5) on the out-of-band comparison.
|
|
plain = self._chain.decrypt(msg4, key=kdf("reg/msg4", self._shared, self._ee))
|
|
try:
|
|
se_ct, tb = REG_MSG4.unpack(plain)
|
|
# Solving se proves to Bob that we hold Alice's static
|
|
se = self._mlkem.decapsulate(se_ct)
|
|
se += self._x25519.exchange(X25519PublicKey.from_public_bytes(self._e_bob))
|
|
except (struct.error, ValueError) as e:
|
|
raise _Burn(f"bad msg4 contents ({e})") from None
|
|
self._se = se
|
|
_check(
|
|
hmac.compare_digest(tb, kdf("reg/confirm-b", self._root())),
|
|
"responder tag mismatch",
|
|
)
|
|
sas = kdf("sas", self._shared, self._ee, self._reveal, length=4)
|
|
msg5 = self._chain.encrypt(
|
|
kdf("reg/confirm-a", self._root()),
|
|
key=kdf("reg/msg5", self._shared, self._ee, self._es),
|
|
)
|
|
return msg5, sas
|
|
|
|
# -- bob ------------------------------------------------------------------
|
|
|
|
def respondB(self, msg1: bytes, opened: bytes) -> bytes:
|
|
# msg2 from an opened msg1 (open_hello did the caller scan); our
|
|
# connection key ships here once, Alice answers every leg with it
|
|
self._chain.advance(msg1)
|
|
try:
|
|
e_alice_m, e_alice_x, commitment = REG_MSG1.unpack(opened)
|
|
self._ephemeral = X25519PrivateKey.generate()
|
|
e_bob = self._ephemeral.public_key().public_bytes_raw()
|
|
# Forward secrecy, made to Alice's connection key only
|
|
ee, ee_ct = MLKEM768PublicKey.from_public_bytes(e_alice_m).encapsulate()
|
|
ee += self._ephemeral.exchange(X25519PublicKey.from_public_bytes(e_alice_x))
|
|
except (struct.error, ValueError) as e:
|
|
raise _Burn(f"bad msg1 contents ({e})") from None
|
|
self._commitment = (e_alice_m, e_alice_x, commitment)
|
|
body = REG_MSG2.pack(
|
|
e_bob,
|
|
ee_ct,
|
|
encrypt(
|
|
REG_MSG2_INNER.pack(self._static_m, self._static_x),
|
|
key=kdf("reg/msg2/inner", ee),
|
|
),
|
|
)
|
|
self._ee = ee
|
|
return self._chain.encrypt(body, key=kdf("reg/msg2", self._shared))
|
|
|
|
def confirmB(self, msg3: bytes) -> tuple[bytes, bytes]:
|
|
# Alice's challenge to our static, her static and her reveal. The
|
|
# caller gates the answer (msg4) on the out-of-band comparison.
|
|
plain = self._chain.decrypt(msg3, key=kdf("reg/msg3", self._shared, self._ee))
|
|
try:
|
|
es_ct, inner = REG_MSG3.unpack(plain)
|
|
e_alice_m, e_alice_x, commitment = self._commitment
|
|
# Solving es proves to Alice that the static key they encapsulated to was ours
|
|
es = self._mlkem.decapsulate(es_ct)
|
|
es += self._x25519.exchange(X25519PublicKey.from_public_bytes(e_alice_x))
|
|
# This enables us to get Alice's static keys and reveal
|
|
s_alice_m, s_alice_x, reveal = REG_MSG3_INNER.unpack(
|
|
decrypt(inner, key=kdf("reg/msg3/inner", es))
|
|
)
|
|
# We send Alice the same test, verifying they have s_alice secrets
|
|
se, se_ct = MLKEM768PublicKey.from_public_bytes(s_alice_m).encapsulate()
|
|
se += self._ephemeral.exchange(X25519PublicKey.from_public_bytes(s_alice_x))
|
|
except (struct.error, ValueError) as e:
|
|
raise _Burn(f"bad msg3 contents ({e})") from None
|
|
_check(
|
|
hmac.compare_digest(
|
|
commitment, kdf("commit", e_alice_m, e_alice_x, reveal)
|
|
),
|
|
"commitment reveal mismatch",
|
|
)
|
|
self._reveal = reveal
|
|
self._peer_pub = (s_alice_m, s_alice_x)
|
|
self._es = es
|
|
self._se = se
|
|
sas = kdf("sas", self._shared, self._ee, reveal, length=4)
|
|
msg4 = self._chain.encrypt(
|
|
REG_MSG4.pack(se_ct, kdf("reg/confirm-b", self._root())),
|
|
key=kdf("reg/msg4", self._shared, self._ee),
|
|
)
|
|
return msg4, sas
|
|
|
|
def finishB(self, msg5: bytes) -> None:
|
|
# Alice's root confirmation: proves she solved our se challenge
|
|
plain = self._chain.decrypt(
|
|
msg5, key=kdf("reg/msg5", self._shared, self._ee, self._es)
|
|
)
|
|
if len(plain) != 32:
|
|
raise _Burn("bad msg5 contents (expected a 32-byte tag)")
|
|
ta = plain
|
|
_check(
|
|
hmac.compare_digest(ta, kdf("reg/confirm-a", self._root())),
|
|
"initiator tag mismatch",
|
|
)
|
|
|
|
# -- completion -------------------------------------------------------------
|
|
|
|
def finish(self, initiator: bool) -> tuple[bytes, bytes, bytes, bytes, tuple]:
|
|
# -> (send, recv, rotated shared, sas, peer pub); the rotation is
|
|
# held by the caller until the store swap
|
|
root = self._root()
|
|
rotated = kdf("rotate", root)
|
|
sas = kdf("sas", self._shared, self._ee, self._reveal, length=4)
|
|
return *_traffic_keys(root, initiator), rotated, sas, self._peer_pub
|
|
|
|
|
|
class Resumption:
|
|
"""Resumption exchange: statics proven by decapsulation, never on
|
|
the wire; all legs hybrid. Alice: startA/respondA; Bob:
|
|
respondB/completeB; finish() -> (send, recv)."""
|
|
|
|
def __init__(self, shared: bytes, idroot: bytes, peer_pub: tuple, ad: bytes):
|
|
self._shared = shared
|
|
self._idroot = idroot
|
|
self._mlkem = _mlkem_private_key(idroot)
|
|
self._x25519 = _x25519_private_key(idroot)
|
|
self._peer_pub = peer_pub # (mlkem public, x25519 public)
|
|
self._chain = _Chain(ad)
|
|
|
|
# -- alice ----------------------------------------------------------------
|
|
|
|
def startA(self) -> bytes:
|
|
# msg1: our connection keys plus a challenge to Bob's static
|
|
self._eph_m, self._ephemeral, e_alice_m, e_alice_x = _fresh_ephemeral()
|
|
peer_m, peer_x = self._peer_pub
|
|
sb, sb_ct = MLKEM768PublicKey.from_public_bytes(peer_m).encapsulate()
|
|
sb += self._ephemeral.exchange(X25519PublicKey.from_public_bytes(peer_x))
|
|
self._sb = sb
|
|
body = RES_MSG1.pack(e_alice_m, e_alice_x, sb_ct)
|
|
return self._chain.encrypt(body, key=kdf("res/msg1", self._shared))
|
|
|
|
def respondA(self, msg2: bytes) -> bytes:
|
|
# Open Bob's reply and solve his three challenges
|
|
plain = self._chain.decrypt(msg2, key=kdf("res/msg2", self._shared, self._sb))
|
|
try:
|
|
e_bob, sa_ct, ee_ct, tb = RES_MSG2.unpack(plain)
|
|
# Solving sa proves to Bob that we hold Alice's static
|
|
sa = self._mlkem.decapsulate(sa_ct)
|
|
sa += self._x25519.exchange(X25519PublicKey.from_public_bytes(e_bob))
|
|
# ee is forward secrecy, made to our connection key
|
|
ee = self._eph_m.decapsulate(ee_ct)
|
|
ee += self._ephemeral.exchange(X25519PublicKey.from_public_bytes(e_bob))
|
|
except (struct.error, ValueError) as e:
|
|
raise _Burn(f"bad msg2 contents ({e})") from None
|
|
_check(
|
|
hmac.compare_digest(tb, kdf("res/tag-b", self._sb, sa, ee)),
|
|
"responder tag mismatch",
|
|
)
|
|
self._sa = sa
|
|
self._ee = ee
|
|
return self._chain.encrypt(
|
|
kdf("res/tag-a", self._sb, sa, ee),
|
|
key=kdf("res/msg3", self._shared, self._sb, sa, ee),
|
|
)
|
|
|
|
# -- bob ------------------------------------------------------------------
|
|
|
|
def respondB(self, msg1: bytes, opened: bytes) -> bytes:
|
|
# Prove our static, challenge Alice's, forward secrecy
|
|
self._chain.advance(msg1)
|
|
try:
|
|
e_alice_m, e_alice_x, sb_ct = RES_MSG1.unpack(opened)
|
|
# Solving sb proves to Alice that we hold Bob's static
|
|
sb = self._mlkem.decapsulate(sb_ct)
|
|
sb += self._x25519.exchange(X25519PublicKey.from_public_bytes(e_alice_x))
|
|
self._ephemeral = X25519PrivateKey.generate()
|
|
e_bob = self._ephemeral.public_key().public_bytes_raw()
|
|
peer_m, peer_x = self._peer_pub
|
|
# We challenge Alice's static the same way
|
|
sa, sa_ct = MLKEM768PublicKey.from_public_bytes(peer_m).encapsulate()
|
|
sa += self._ephemeral.exchange(X25519PublicKey.from_public_bytes(peer_x))
|
|
# And make forward secrecy to her connection key
|
|
ee, ee_ct = MLKEM768PublicKey.from_public_bytes(e_alice_m).encapsulate()
|
|
ee += self._ephemeral.exchange(X25519PublicKey.from_public_bytes(e_alice_x))
|
|
except (struct.error, ValueError) as e:
|
|
raise _Burn(f"bad msg1 contents ({e})") from None
|
|
self._sb = sb
|
|
self._sa = sa
|
|
self._ee = ee
|
|
body = RES_MSG2.pack(e_bob, sa_ct, ee_ct, kdf("res/tag-b", sb, sa, ee))
|
|
return self._chain.encrypt(body, key=kdf("res/msg2", self._shared, sb))
|
|
|
|
def completeB(self, msg3: bytes) -> None:
|
|
# Alice answers with just her tag
|
|
ta = self._chain.decrypt(
|
|
msg3, key=kdf("res/msg3", self._shared, self._sb, self._sa, self._ee)
|
|
)
|
|
_check(
|
|
hmac.compare_digest(ta, kdf("res/tag-a", self._sb, self._sa, self._ee)),
|
|
"initiator tag mismatch",
|
|
)
|
|
|
|
# -- completion -------------------------------------------------------------
|
|
|
|
def finish(self, initiator: bool) -> tuple[bytes, bytes]:
|
|
# -> (send, recv); a resumption rotates nothing
|
|
key = kdf("res/session-key", self._shared, self._sb, self._sa, self._ee)
|
|
return _traffic_keys(key, initiator)
|