diff options
Diffstat (limited to 'packages/meshbay-common/src/meshbay_common')
| -rw-r--r-- | packages/meshbay-common/src/meshbay_common/ratchet.py | 311 |
1 files changed, 0 insertions, 311 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) |