diff options
| -rw-r--r-- | packages/meshbay-common/src/meshbay_common/ratchet.py | 311 | ||||
| -rw-r--r-- | packages/meshbay-common/tests/test_ratchet.py | 167 |
2 files changed, 478 insertions, 0 deletions
diff --git a/packages/meshbay-common/src/meshbay_common/ratchet.py b/packages/meshbay-common/src/meshbay_common/ratchet.py new file mode 100644 index 0000000..1750225 --- /dev/null +++ b/packages/meshbay-common/src/meshbay_common/ratchet.py @@ -0,0 +1,311 @@ +""" +MeshBay — Double Ratchet Algorithm implementation. + +Based on the Signal Protocol specification: + https://signal.org/docs/specifications/doubleratchet/ + +Provides forward secrecy and break-in recovery for group messaging. +Each message is encrypted with a unique key derived from the ratchet state. +Compromise of the current state does not reveal past message keys. + +Key components: + - DH Ratchet: rotates Diffie-Hellman keys to achieve break-in recovery + - Symmetric Ratchet: derives unique per-message keys from a chain key + - KDF functions: HKDF-based key derivation following the Signal spec + +Usage: + # Initialise from a shared secret (e.g. from X3DH or GEK) + alice_state = RatchetState.init_sender(shared_secret, bob_public_key) + bob_state = RatchetState.init_receiver(shared_secret, bob_private_key) + + # Alice sends + header, ciphertext = alice_state.encrypt(b"Hello Bob") + + # Bob receives + plaintext = bob_state.decrypt(header, ciphertext) + assert plaintext == b"Hello Bob" +""" + +import os +import struct +from dataclasses import dataclass, field +from typing import Optional + +from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey, X25519PublicKey +from cryptography.hazmat.primitives.ciphers.aead import AESGCM +from cryptography.hazmat.primitives.kdf.hkdf import HKDF +from cryptography.hazmat.primitives import hashes, serialization + +# KDF info strings (stable identifiers) +_KDF_RK_INFO = b"meshbay:ratchet:root:v1" +_KDF_CK_INFO = b"meshbay:ratchet:chain:v1" +_KDF_MSG_INFO = b"meshbay:ratchet:msg:v1" + +# Maximum out-of-order messages stored per ratchet step +MAX_SKIP = 1000 + + +# ── Key helpers ─────────────────────────────────────────────────────────────── + +def _dh_generate() -> X25519PrivateKey: + return X25519PrivateKey.generate() + +def _dh_pub_bytes(sk: X25519PrivateKey) -> bytes: + return sk.public_key().public_bytes( + serialization.Encoding.Raw, serialization.PublicFormat.Raw) + +def _dh_pub_from_bytes(b: bytes) -> X25519PublicKey: + return X25519PublicKey.from_public_bytes(b) + +def _dh(sk: X25519PrivateKey, pk: X25519PublicKey) -> bytes: + return sk.exchange(pk) + +def _kdf_rk(root_key: bytes, dh_out: bytes) -> tuple[bytes, bytes]: + """KDF_RK: derive new root key + chain key from root key + DH output.""" + material = HKDF( + algorithm=hashes.SHA256(), length=64, + salt=root_key, info=_KDF_RK_INFO, + ).derive(dh_out) + return material[:32], material[32:] # new_rk, ck + +def _kdf_ck(chain_key: bytes) -> tuple[bytes, bytes]: + """KDF_CK: advance chain key, derive message key.""" + material = HKDF( + algorithm=hashes.SHA256(), length=64, + salt=chain_key, info=_KDF_CK_INFO, + ).derive(b"\x01") + return material[:32], material[32:] # new_ck, mk + +def _kdf_msg(msg_key: bytes) -> tuple[bytes, bytes, bytes]: + """Expand msg_key into encryption_key, auth_key, iv (96-bit nonce).""" + material = HKDF( + algorithm=hashes.SHA256(), length=80, + salt=b"\x00" * 32, info=_KDF_MSG_INFO, + ).derive(msg_key) + return material[:32], material[32:64], material[64:] # enc_key, auth_key, iv + +def _encrypt(mk: bytes, plaintext: bytes, ad: bytes) -> bytes: + enc_key, _, iv = _kdf_msg(mk) + return AESGCM(enc_key).encrypt(iv, plaintext, ad) + +def _decrypt(mk: bytes, ciphertext: bytes, ad: bytes) -> bytes: + enc_key, _, iv = _kdf_msg(mk) + return AESGCM(enc_key).decrypt(iv, ciphertext, ad) + + +# ── Message header ──────────────────────────────────────────────────────────── + +@dataclass +class MessageHeader: + """ + Wire header for a ratchet-encrypted message. + dh_pub: sender's current DH public key (32 bytes raw) + prev_chain_n: number of messages in previous sending chain + msg_num: message number in current sending chain + """ + dh_pub: bytes + prev_chain_n: int + msg_num: int + + def encode(self) -> bytes: + return self.dh_pub + struct.pack(">II", self.prev_chain_n, self.msg_num) + + @classmethod + def decode(cls, data: bytes) -> "MessageHeader": + dh_pub = data[:32] + prev_chain_n, msg_num = struct.unpack(">II", data[32:40]) + return cls(dh_pub=dh_pub, prev_chain_n=prev_chain_n, msg_num=msg_num) + + @property + def encoded_len(self) -> int: + return 40 # 32 (dh) + 4 + 4 + + +# ── Ratchet state ───────────────────────────────────────────────────────────── + +@dataclass +class RatchetState: + """ + Full Double Ratchet state for one participant. + + Fields match Signal spec (Appendix C): + DHs: our sending DH key pair + DHr: remote's DH public key (None until first message received) + RK: 32-byte root key + CKs: sending chain key (None until first send) + CKr: receiving chain key (None until first receive) + Ns: number of messages sent on current sending chain + Nr: number of messages received on current receiving chain + PN: number of messages sent in previous sending chain + MKSKIP: skipped message keys {(dh_pub_bytes, msg_num) → mk} + """ + DHs: X25519PrivateKey + DHr: Optional[X25519PublicKey] + RK: bytes + CKs: Optional[bytes] + CKr: Optional[bytes] + Ns: int = 0 + Nr: int = 0 + PN: int = 0 + MKSKIP: dict = field(default_factory=dict) + + # ── Initialisation ──────────────────────────────────────────────────────── + + @classmethod + def init_sender( + cls, shared_secret: bytes, remote_public_key: bytes + ) -> "RatchetState": + """ + Initialise state as the message sender. + shared_secret: pre-shared key from X3DH or GEK + remote_public_key: recipient's initial X25519 public key (raw 32 bytes) + """ + dhs = _dh_generate() + dhr = _dh_pub_from_bytes(remote_public_key) + rk, cks = _kdf_rk(shared_secret, _dh(dhs, dhr)) + return cls(DHs=dhs, DHr=dhr, RK=rk, CKs=cks, CKr=None) + + @classmethod + def init_receiver( + cls, shared_secret: bytes, own_private_key: X25519PrivateKey + ) -> "RatchetState": + """ + Initialise state as the message receiver. + shared_secret: same pre-shared key used by sender + own_private_key: the X25519 private key whose public key was given to sender + """ + return cls( + DHs=own_private_key, + DHr=None, + RK=shared_secret, + CKs=None, + CKr=None, + ) + + # ── Encryption ──────────────────────────────────────────────────────────── + + def encrypt(self, plaintext: bytes, associated_data: bytes = b"") -> tuple[MessageHeader, bytes]: + """Encrypt a message. Returns (header, ciphertext).""" + assert self.CKs is not None, "Not initialised as sender" + self.CKs, mk = _kdf_ck(self.CKs) + header = MessageHeader( + dh_pub=_dh_pub_bytes(self.DHs), + prev_chain_n=self.PN, + msg_num=self.Ns, + ) + self.Ns += 1 + ct = _encrypt(mk, plaintext, associated_data + header.encode()) + return header, ct + + # ── Decryption ──────────────────────────────────────────────────────────── + + def decrypt(self, header: MessageHeader, ciphertext: bytes, + associated_data: bytes = b"") -> bytes: + """Decrypt a message. Handles out-of-order delivery via MKSKIP.""" + ad = associated_data + header.encode() + + # Check skipped keys first + skip_key = (header.dh_pub, header.msg_num) + if skip_key in self.MKSKIP: + mk = self.MKSKIP.pop(skip_key) + return _decrypt(mk, ciphertext, ad) + + # DH ratchet step if new DH key received + if self.DHr is None or header.dh_pub != _dh_pub_bytes_from_pk(self.DHr): + self._skip_message_keys(header.prev_chain_n) + self._dh_ratchet(header) + + # Advance receiving chain + self._skip_message_keys(header.msg_num) + assert self.CKr is not None + self.CKr, mk = _kdf_ck(self.CKr) + self.Nr += 1 + return _decrypt(mk, ciphertext, ad) + + def _skip_message_keys(self, until: int) -> None: + """Store skipped message keys for out-of-order delivery.""" + if self.Nr + MAX_SKIP < until: + raise ValueError(f"Too many skipped messages: {until - self.Nr}") + if self.CKr is None: + return + while self.Nr < until: + self.CKr, mk = _kdf_ck(self.CKr) + self.MKSKIP[(bytes(_dh_pub_bytes_from_pk(self.DHr)), self.Nr)] = mk + self.Nr += 1 + + def _dh_ratchet(self, header: MessageHeader) -> None: + """Perform a DH ratchet step on receiving a new remote DH key.""" + self.PN = self.Ns + self.Ns = 0 + self.Nr = 0 + self.DHr = _dh_pub_from_bytes(header.dh_pub) + self.RK, self.CKr = _kdf_rk(self.RK, _dh(self.DHs, self.DHr)) + self.DHs = _dh_generate() + self.RK, self.CKs = _kdf_rk(self.RK, _dh(self.DHs, self.DHr)) + + +def _dh_pub_bytes_from_pk(pk: X25519PublicKey | None) -> bytes: + if pk is None: + return b"\x00" * 32 + return pk.public_bytes(serialization.Encoding.Raw, serialization.PublicFormat.Raw) + + +# ── Group chat helpers ──────────────────────────────────────────────────────── + +@dataclass +class ChatMessage: + """A ratchet-encrypted group chat message.""" + sender_id: str + header_enc: bytes # encoded MessageHeader + ciphertext: bytes + timestamp: int # unix timestamp + msg_id: str # uuid4 + + def to_dict(self) -> dict: + import base64 + return { + "sender_id": self.sender_id, + "header_b64": base64.b64encode(self.header_enc).decode(), + "ct_b64": base64.b64encode(self.ciphertext).decode(), + "timestamp": self.timestamp, + "msg_id": self.msg_id, + } + + @classmethod + def from_dict(cls, d: dict) -> "ChatMessage": + import base64 + return cls( + sender_id=d["sender_id"], + header_enc=base64.b64decode(d["header_b64"]), + ciphertext=base64.b64decode(d["ct_b64"]), + timestamp=d["timestamp"], + msg_id=d["msg_id"], + ) + + +def encrypt_chat_message( + state: RatchetState, + sender_id: str, + plaintext: str | bytes, +) -> ChatMessage: + """Encrypt a chat message using the Double Ratchet state.""" + import time, uuid + if isinstance(plaintext, str): + plaintext = plaintext.encode() + header, ct = state.encrypt(plaintext) + return ChatMessage( + sender_id=sender_id, + header_enc=header.encode(), + ciphertext=ct, + timestamp=int(time.time()), + msg_id=str(uuid.uuid4()), + ) + + +def decrypt_chat_message( + state: RatchetState, + msg: ChatMessage, +) -> bytes: + """Decrypt a chat message using the Double Ratchet state.""" + header = MessageHeader.decode(msg.header_enc) + return state.decrypt(header, msg.ciphertext) diff --git a/packages/meshbay-common/tests/test_ratchet.py b/packages/meshbay-common/tests/test_ratchet.py new file mode 100644 index 0000000..66d2039 --- /dev/null +++ b/packages/meshbay-common/tests/test_ratchet.py @@ -0,0 +1,167 @@ +"""Tests for the Double Ratchet implementation.""" + +import os +import pytest +from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey + +from meshbay_common.ratchet import ( + RatchetState, + MessageHeader, + ChatMessage, + encrypt_chat_message, + decrypt_chat_message, +) + + +@pytest.fixture +def shared_secret(): + return os.urandom(32) + + +@pytest.fixture +def bob_key(): + return X25519PrivateKey.generate() + + +@pytest.fixture +def alice_bob(shared_secret, bob_key): + from cryptography.hazmat.primitives import serialization + bob_pub_raw = bob_key.public_key().public_bytes( + serialization.Encoding.Raw, serialization.PublicFormat.Raw) + alice = RatchetState.init_sender(shared_secret, bob_pub_raw) + bob = RatchetState.init_receiver(shared_secret, bob_key) + return alice, bob + + +def test_basic_send_receive(alice_bob): + alice, bob = alice_bob + header, ct = alice.encrypt(b"Hello Bob!") + plaintext = bob.decrypt(header, ct) + assert plaintext == b"Hello Bob!" + + +def test_multiple_messages(alice_bob): + alice, bob = alice_bob + messages = [b"msg 1", b"msg 2", b"msg 3", b"msg 4", b"msg 5"] + ciphertexts = [alice.encrypt(m) for m in messages] + for i, (h, ct) in enumerate(ciphertexts): + assert bob.decrypt(h, ct) == messages[i] + + +def test_bidirectional(alice_bob): + """Both sides can send and receive.""" + alice, bob = alice_bob + + # Alice → Bob + h, ct = alice.encrypt(b"Hi from Alice") + assert bob.decrypt(h, ct) == b"Hi from Alice" + + # Bob → Alice (triggers DH ratchet on Alice) + h, ct = bob.encrypt(b"Hi from Bob") + assert alice.decrypt(h, ct) == b"Hi from Bob" + + # Alice → Bob again (new chain) + h, ct = alice.encrypt(b"Second message from Alice") + assert bob.decrypt(h, ct) == b"Second message from Alice" + + +def test_forward_secrecy(alice_bob): + """Encrypting msg N+1 erases the key for msg N.""" + alice, bob = alice_bob + h1, ct1 = alice.encrypt(b"message 1") + h2, ct2 = alice.encrypt(b"message 2") + + assert bob.decrypt(h1, ct1) == b"message 1" + assert bob.decrypt(h2, ct2) == b"message 2" + + # Keys are consumed — cannot replay + with pytest.raises(Exception): + bob.decrypt(h1, ct1) + + +def test_out_of_order_delivery(alice_bob): + """Messages arriving out of order should still decrypt correctly.""" + alice, bob = alice_bob + h1, ct1 = alice.encrypt(b"first") + h2, ct2 = alice.encrypt(b"second") + h3, ct3 = alice.encrypt(b"third") + + # Deliver in reverse order + assert bob.decrypt(h3, ct3) == b"third" + assert bob.decrypt(h2, ct2) == b"second" + assert bob.decrypt(h1, ct1) == b"first" + + +def test_associated_data(alice_bob): + alice, bob = alice_bob + ad = b"sender:alice;group:test-group" + h, ct = alice.encrypt(b"secret", associated_data=ad) + assert bob.decrypt(h, ct, associated_data=ad) == b"secret" + + # Wrong AD fails authentication + with pytest.raises(Exception): + bob.decrypt(h, ct, associated_data=b"wrong-ad") + + +def test_header_encode_decode(): + from cryptography.hazmat.primitives import serialization + sk = X25519PrivateKey.generate() + pub = sk.public_key().public_bytes( + serialization.Encoding.Raw, serialization.PublicFormat.Raw) + h = MessageHeader(dh_pub=pub, prev_chain_n=5, msg_num=12) + decoded = MessageHeader.decode(h.encode()) + assert decoded.dh_pub == pub + assert decoded.prev_chain_n == 5 + assert decoded.msg_num == 12 + + +def test_many_messages_stress(alice_bob): + """100 messages without DH ratchet — verifies chain key stability.""" + alice, bob = alice_bob + for i in range(100): + h, ct = alice.encrypt(f"message {i}".encode()) + assert bob.decrypt(h, ct) == f"message {i}".encode() + + +def test_encrypt_decrypt_chat_message(alice_bob, shared_secret, bob_key): + alice, bob = alice_bob + msg = encrypt_chat_message(alice, "alice_user_id", "Hello group!") + assert msg.sender_id == "alice_user_id" + assert msg.msg_id != "" + + plaintext = decrypt_chat_message(bob, msg) + assert plaintext == b"Hello group!" + + +def test_chat_message_serialization(alice_bob): + alice, bob = alice_bob + msg = encrypt_chat_message(alice, "alice", "Serialize me") + d = msg.to_dict() + restored = ChatMessage.from_dict(d) + assert decrypt_chat_message(bob, restored) == b"Serialize me" + + +def test_break_in_recovery(shared_secret, bob_key): + """ + Compromise of state at message N does not reveal keys for messages > N. + After a DH ratchet step, new keys are independent of the compromised state. + """ + from cryptography.hazmat.primitives import serialization + bob_pub = bob_key.public_key().public_bytes( + serialization.Encoding.Raw, serialization.PublicFormat.Raw) + + alice = RatchetState.init_sender(shared_secret, bob_pub) + bob = RatchetState.init_receiver(shared_secret, bob_key) + + # Exchange some messages + h, ct = alice.encrypt(b"pre-compromise msg") + bob.decrypt(h, ct) + + # Bob replies (triggers DH ratchet — new keys independent of above) + h, ct = bob.encrypt(b"bob reply triggers ratchet") + alice.decrypt(h, ct) + + # Now Alice's state has fresh keys + h, ct = alice.encrypt(b"post-ratchet msg") + plaintext = bob.decrypt(h, ct) + assert plaintext == b"post-ratchet msg" |