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, 0 insertions, 478 deletions
diff --git a/packages/meshbay-common/src/meshbay_common/ratchet.py b/packages/meshbay-common/src/meshbay_common/ratchet.py deleted file mode 100644 index 1750225..0000000 --- a/packages/meshbay-common/src/meshbay_common/ratchet.py +++ /dev/null @@ -1,311 +0,0 @@ -""" -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 deleted file mode 100644 index 66d2039..0000000 --- a/packages/meshbay-common/tests/test_ratchet.py +++ /dev/null @@ -1,167 +0,0 @@ -"""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" |