1 files changed, 167 insertions, 0 deletions
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"