368 lines
15 KiB
Python
368 lines
15 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 import mlkem, x25519
|
|
|
|
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)
|
|
_PUB = struct.Struct("1184s32s") # hybrid public key: 1216
|
|
_CAP = struct.Struct("1088s32s") # hybrid capsule: mlkem ct | fresh x25519 pub
|
|
_PUB_CAP = struct.Struct("1216s1120s") # hybrid public key | hybrid capsule
|
|
_TAG = struct.Struct("32s") # bare kdf tag
|
|
|
|
# handshake message bodies
|
|
_MSG1_REGISTER = struct.Struct("1216s32s") # e_alice | commitment
|
|
_STATIC = struct.Struct("1216s") # bare hybrid static public key
|
|
_MSG3_REGISTER_INNER = struct.Struct("1216s32s") # s_alice | reveal
|
|
_MSG4_REGISTER = struct.Struct("1120s32s") # se_cap | tb
|
|
_MSG2_RESUME = struct.Struct("1120s1120s32s") # sa_cap | ee_cap | tb
|
|
# msg2 = ee_cap | inner box(s_bob); msg3 = es_cap | inner box(s_alice | reveal);
|
|
# msg5 = _TAG (ta); msg1 resumption = _PUB_CAP (e_alice | sb_cap);
|
|
# 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, operating on raw byte buffers only
|
|
#
|
|
# A 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.
|
|
|
|
|
|
def _mlkem_private_key(idroot: bytes) -> mlkem.MLKEM768PrivateKey:
|
|
return mlkem.MLKEM768PrivateKey.from_seed_bytes(
|
|
kdf("idroot/mlkem", idroot, length=64)
|
|
)
|
|
|
|
|
|
def _x25519_private_key(idroot: bytes) -> x25519.X25519PrivateKey:
|
|
return x25519.X25519PrivateKey.from_private_bytes(kdf("idroot/x25519", idroot))
|
|
|
|
|
|
def _public_from_idroot(idroot: bytes) -> bytes:
|
|
return _PUB.pack(
|
|
_mlkem_private_key(idroot).public_key().public_bytes_raw(),
|
|
_x25519_private_key(idroot).public_key().public_bytes_raw(),
|
|
)
|
|
|
|
|
|
def _encapsulate(peer_public: bytes) -> tuple[bytes, bytes]:
|
|
# -> (shared_secret, capsule), each the concat of both schemes
|
|
mk_raw, xk_raw = _PUB.unpack(peer_public)
|
|
m_ss, m_ct = mlkem.MLKEM768PublicKey.from_public_bytes(mk_raw).encapsulate()
|
|
x_ek = x25519.X25519PrivateKey.generate()
|
|
try:
|
|
x_ss = x_ek.exchange(x25519.X25519PublicKey.from_public_bytes(xk_raw))
|
|
except ValueError:
|
|
raise ValueError("non-contributory X25519 public key") from None
|
|
return m_ss + x_ss, _CAP.pack(m_ct, x_ek.public_key().public_bytes_raw())
|
|
|
|
|
|
def _decapsulate(idroot: bytes, cap: bytes) -> bytes:
|
|
# ML-KEM has FIPS 203 implicit rejection (never fails; wrong key ->
|
|
# garbage caught by the kdf tags); X25519 does not, so a low-order
|
|
# result is rejected explicitly
|
|
m_ct, x_ct = _CAP.unpack(cap)
|
|
m_ss = _mlkem_private_key(idroot).decapsulate(m_ct)
|
|
try:
|
|
x_ss = _x25519_private_key(idroot).exchange(
|
|
x25519.X25519PublicKey.from_public_bytes(x_ct)
|
|
)
|
|
except ValueError:
|
|
raise ValueError("non-contributory X25519 capsule") from None
|
|
return m_ss + x_ss
|
|
|
|
|
|
def _fresh_ephemeral() -> tuple[bytes, bytes]:
|
|
# (root, public) for a fresh per-connection ephemeral key
|
|
root = random(32)
|
|
return root, _public_from_idroot(root)
|
|
|
|
|
|
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 the same
|
|
hybrid KEM operation -- ee_cap for forward secrecy, es_cap/se_cap
|
|
challenging the statics passed late (s_bob inside msg2's inner box,
|
|
s_alice inside msg3's). finish() -> (send, recv, rotated shared,
|
|
sas, peer pub)."""
|
|
|
|
def __init__(self, shared: bytes, idroot: bytes, ad: bytes):
|
|
self._shared = shared
|
|
self._idroot = idroot
|
|
self._static_pub = _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: the ephemeral key, committed under a reveal hidden until msg3
|
|
self._ephemeral_root, e_alice = _fresh_ephemeral()
|
|
self._reveal = random(32)
|
|
body = _MSG1_REGISTER.pack(e_alice, kdf("commit", e_alice, self._reveal))
|
|
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 the static inside its inner
|
|
# box, and challenge it
|
|
plain = self._chain.decrypt(msg2, key=kdf("reg/msg2", self._shared))
|
|
try:
|
|
ee_cap, inner = plain[: _CAP.size], plain[_CAP.size :]
|
|
ee = _decapsulate(self._ephemeral_root, ee_cap)
|
|
(s_bob,) = _STATIC.unpack(decrypt(inner, key=kdf("reg/msg2/inner", ee)))
|
|
es, es_cap = _encapsulate(s_bob) # challenge: bob's static
|
|
except (struct.error, ValueError) as e:
|
|
raise _Burn(f"bad msg2 contents ({e})") from None
|
|
self._peer_pub = s_bob
|
|
self._ee = ee
|
|
self._es = es
|
|
return self._chain.encrypt(
|
|
es_cap
|
|
+ encrypt(
|
|
_MSG3_REGISTER_INNER.pack(self._static_pub, 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_cap, tb = _MSG4_REGISTER.unpack(plain)
|
|
se = _decapsulate(self._idroot, se_cap)
|
|
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)
|
|
self._chain.advance(msg1)
|
|
try:
|
|
e_alice, commitment = _MSG1_REGISTER.unpack(opened)
|
|
ee, ee_cap = _encapsulate(e_alice) # forward secrecy
|
|
except (struct.error, ValueError) as e:
|
|
raise _Burn(f"bad msg1 contents ({e})") from None
|
|
self._commitment = (e_alice, commitment)
|
|
body = ee_cap + encrypt(
|
|
_STATIC.pack(self._static_pub), 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_cap, inner = plain[: _CAP.size], plain[_CAP.size :]
|
|
es = _decapsulate(self._idroot, es_cap) # proves our static
|
|
s_alice, reveal = _MSG3_REGISTER_INNER.unpack(
|
|
decrypt(inner, key=kdf("reg/msg3/inner", es))
|
|
)
|
|
se, se_cap = _encapsulate(s_alice)
|
|
except (struct.error, ValueError) as e:
|
|
raise _Burn(f"bad msg3 contents ({e})") from None
|
|
e_alice, commitment = self._commitment
|
|
_check(
|
|
hmac.compare_digest(commitment, kdf("commit", e_alice, reveal)),
|
|
"commitment reveal mismatch",
|
|
)
|
|
self._reveal = reveal
|
|
self._peer_pub = s_alice
|
|
self._es = es
|
|
self._se = se
|
|
sas = kdf("sas", self._shared, self._ee, reveal, length=4)
|
|
msg4 = self._chain.encrypt(
|
|
_MSG4_REGISTER.pack(se_cap, 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 decapsulated se_cap
|
|
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, bytes]:
|
|
# -> (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: bytes, ad: bytes):
|
|
self._shared = shared
|
|
self._idroot = idroot
|
|
self._peer_pub = peer_pub # proven by decapsulation
|
|
self._chain = _Chain(ad)
|
|
|
|
# -- alice ----------------------------------------------------------------
|
|
|
|
def startA(self) -> bytes:
|
|
# msg1: the ephemeral key plus an encapsulation to bob's static
|
|
# (the responder ships no ephemeral of his own: no leg uses it)
|
|
self._ephemeral_root, e_alice = _fresh_ephemeral()
|
|
self._sb, sb_cap = _encapsulate(self._peer_pub)
|
|
body = _PUB_CAP.pack(e_alice, sb_cap)
|
|
return self._chain.encrypt(body, key=kdf("res/msg1", self._shared))
|
|
|
|
def respondA(self, msg2: bytes) -> bytes:
|
|
# open bob's reply, decapsulate all three legs
|
|
plain = self._chain.decrypt(msg2, key=kdf("res/msg2", self._shared, self._sb))
|
|
try:
|
|
sa_cap, ee_cap, tb = _MSG2_RESUME.unpack(plain)
|
|
sa = _decapsulate(self._idroot, sa_cap) # proves our static
|
|
ee = _decapsulate(self._ephemeral_root, ee_cap) # forward secrecy
|
|
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, prove alice's, forward secrecy
|
|
self._chain.advance(msg1)
|
|
try:
|
|
e_alice, sb_cap = _PUB_CAP.unpack(opened)
|
|
sb = _decapsulate(self._idroot, sb_cap) # proves our static
|
|
sa, sa_cap = _encapsulate(self._peer_pub) # proves peer's static
|
|
ee, ee_cap = _encapsulate(e_alice) # forward secrecy
|
|
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(sa_cap, ee_cap, 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)
|