Files
quantumtunnel/quantumpipe/handshake.py
T
LeoVasanko f035859fd2 Split hybrid public keys into separate wire fields
The packed 1216-byte hybrid public had no consumer that wanted it
whole -- everything unpacked immediately. Message layouts and the
credential record now carry peer_m | peer_x as explicit fields; _PUB
and the unpack-at-use idiom are gone. credential_id values change with
the record layout. Step comments in the exchange methods explain what
each solve/challenge proves.
2026-10-07 17:59:13 +00:00

408 lines
17 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)
_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)