"""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) _TAG = struct.Struct("32s") # Bare kdf tag # Handshake message bodies; box(N) is a chab3 box of N plaintext bytes (N+28) _MSG1_REGISTER = struct.Struct("1184s32s32s") # e_alice_m | e_alice_x | commitment _MSG2_REGISTER = struct.Struct( "32s1088s1244s" ) # e_bob | ee_ct | box(s_bob_m | s_bob_x) _MSG2_REGISTER_INNER = struct.Struct("1184s32s") # s_bob_m | s_bob_x _MSG3_REGISTER = struct.Struct("1088s1276s") # es_ct | box(s_alice_m | s_alice_x | r) _MSG3_REGISTER_INNER = struct.Struct("1184s32s32s") # s_alice_m | s_alice_x | reveal _MSG4_REGISTER = struct.Struct("1088s32s") # se_ct | tb _MSG1_RESUME = struct.Struct("1184s32s1088s") # e_alice_m | e_alice_x | sb_ct _MSG2_RESUME = struct.Struct("32s1088s1088s32s") # e_bob | sa_ct | ee_ct | tb # msg5 = _TAG (ta); msg3 resumption = _TAG (ta) # kdf info strings: boxes are namespaced per flow and message (reg/msg1, # res/msg2, reg/msg2/inner, ...) so the two exchanges are disjoint; other # derivations are bare functional labels, shared between the flows when # the inputs have the same format (a2b/b2a, idroot/*, cred/id). # -------------------------------------------------------------------------- # Hybrid KEM: ML-KEM-768 || X25519 # # A static private key is a single 32-byte idroot secret: each scheme's # key is derived from it with one domain-separated kdf call, so the two # subkeys are independent PRF outputs and the idroot is the only thing # held (and the only thing the credential stores). Ephemeral keys are # generated directly -- nothing of them persists. # # Each party ships one X25519 ephemeral public per connection (Alice's # inside e_alice in msg1, Bob's as e_bob in msg2); every X25519 leg is a # plain DH against those shipped keys. ML-KEM legs ship as bare ML-KEM # ciphertexts. 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 = _MSG1_REGISTER.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 = _MSG2_REGISTER.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 = _MSG2_REGISTER_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( _MSG3_REGISTER_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 = _MSG4_REGISTER.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 = _MSG1_REGISTER.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 = _MSG2_REGISTER.pack( e_bob, ee_ct, encrypt( _MSG2_REGISTER_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 = _MSG3_REGISTER.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 = _MSG3_REGISTER_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( _MSG4_REGISTER.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) ) try: (ta,) = _TAG.unpack(plain) except struct.error as e: raise _Burn(f"bad msg5 contents ({e})") from None _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 = _MSG1_RESUME.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 = _MSG2_RESUME.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 = _MSG1_RESUME.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 = _MSG2_RESUME.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 plain = self._chain.decrypt( msg3, key=kdf("res/msg3", self._shared, self._sb, self._sa, self._ee) ) try: (ta,) = _TAG.unpack(plain) except struct.error as e: raise _Burn(f"bad msg3 contents ({e})") from None _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)