1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
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"
|