A Python port of the Invisible Internet Project (I2P)
1"""Tests for HybridDHState -- X25519 + ML-KEM hybrid Noise handshake.
2
3TDD: tests written before implementation.
4"""
5from __future__ import annotations
6
7import hashlib
8
9import pytest
10
11from i2p_crypto.mlkem import MLKEMVariant, is_available as mlkem_is_available
12
13# Skip all tests if pqcrypto is not installed
14pytestmark = pytest.mark.skipif(
15 not mlkem_is_available(), reason="pqcrypto not installed"
16)
17
18
19from i2p_crypto.noise import HybridDHState
20
21
22# ---------- key sizes ----------
23
24
25class TestHybridDHKeySizes:
26 def test_hybrid_dh_key_sizes_512(self) -> None:
27 h = HybridDHState(MLKEMVariant.ML_KEM_512)
28 assert h.public_key_len == 32 + MLKEMVariant.ML_KEM_512.public_key_len
29 assert h.ciphertext_len == 32 + MLKEMVariant.ML_KEM_512.ciphertext_len
30
31 def test_hybrid_dh_key_sizes_768(self) -> None:
32 h = HybridDHState(MLKEMVariant.ML_KEM_768)
33 assert h.public_key_len == 32 + MLKEMVariant.ML_KEM_768.public_key_len
34 assert h.ciphertext_len == 32 + MLKEMVariant.ML_KEM_768.ciphertext_len
35
36 def test_hybrid_dh_key_sizes_1024(self) -> None:
37 h = HybridDHState(MLKEMVariant.ML_KEM_1024)
38 assert h.public_key_len == 32 + MLKEMVariant.ML_KEM_1024.public_key_len
39 assert h.ciphertext_len == 32 + MLKEMVariant.ML_KEM_1024.ciphertext_len
40
41
42# ---------- roundtrip ----------
43
44
45class TestHybridRoundtrip:
46 def test_hybrid_roundtrip(self) -> None:
47 """Alice generates keypair, Bob encapsulates, Alice decapsulates -- same secret."""
48 alice = HybridDHState(MLKEMVariant.ML_KEM_768)
49 alice.generate_keypair()
50
51 bob = HybridDHState(MLKEMVariant.ML_KEM_768)
52 bob.set_remote_public_key(alice.get_public_key())
53
54 response, bob_secret = bob.encapsulate()
55 alice_secret = alice.decapsulate(response)
56
57 assert alice_secret == bob_secret
58 assert len(alice_secret) == 32 # SHA-256 output
59
60 def test_hybrid_different_variants(self) -> None:
61 """All three ML-KEM variants produce valid handshakes."""
62 for variant in MLKEMVariant:
63 alice = HybridDHState(variant)
64 alice.generate_keypair()
65
66 bob = HybridDHState(variant)
67 bob.set_remote_public_key(alice.get_public_key())
68
69 response, bob_secret = bob.encapsulate()
70 alice_secret = alice.decapsulate(response)
71
72 assert alice_secret == bob_secret, f"Failed for {variant}"
73
74 def test_hybrid_different_keypairs_different_secrets(self) -> None:
75 """Two independent handshakes produce different shared secrets."""
76 alice1 = HybridDHState(MLKEMVariant.ML_KEM_768)
77 alice1.generate_keypair()
78 bob1 = HybridDHState(MLKEMVariant.ML_KEM_768)
79 bob1.set_remote_public_key(alice1.get_public_key())
80 _, secret1 = bob1.encapsulate()
81
82 alice2 = HybridDHState(MLKEMVariant.ML_KEM_768)
83 alice2.generate_keypair()
84 bob2 = HybridDHState(MLKEMVariant.ML_KEM_768)
85 bob2.set_remote_public_key(alice2.get_public_key())
86 _, secret2 = bob2.encapsulate()
87
88 assert secret1 != secret2
89
90
91# ---------- format verification ----------
92
93
94class TestHybridFormats:
95 def test_public_key_format(self) -> None:
96 """get_public_key() returns x25519_pub || mlkem_pub concatenation."""
97 h = HybridDHState(MLKEMVariant.ML_KEM_768)
98 h.generate_keypair()
99
100 pub = h.get_public_key()
101 assert len(pub) == h.public_key_len
102
103 # First 32 bytes are X25519 public key
104 x25519_pub = pub[:32]
105 assert len(x25519_pub) == 32
106
107 # Remaining bytes are ML-KEM public key
108 mlkem_pub = pub[32:]
109 assert len(mlkem_pub) == MLKEMVariant.ML_KEM_768.public_key_len
110
111 def test_encapsulate_response_format(self) -> None:
112 """encapsulate() returns (x25519_pub || mlkem_ct, secret)."""
113 alice = HybridDHState(MLKEMVariant.ML_KEM_768)
114 alice.generate_keypair()
115
116 bob = HybridDHState(MLKEMVariant.ML_KEM_768)
117 bob.set_remote_public_key(alice.get_public_key())
118
119 response, _ = bob.encapsulate()
120 assert len(response) == bob.ciphertext_len
121
122 # First 32 bytes are Bob's X25519 public key
123 bob_x25519_pub = response[:32]
124 assert len(bob_x25519_pub) == 32
125
126 # Remaining bytes are ML-KEM ciphertext
127 mlkem_ct = response[32:]
128 assert len(mlkem_ct) == MLKEMVariant.ML_KEM_768.ciphertext_len
129
130
131# ---------- secret derivation ----------
132
133
134class TestHybridSecretDerivation:
135 def test_hybrid_secret_is_sha256(self) -> None:
136 """Verify shared secret = SHA-256(x25519_ss || mlkem_ss).
137
138 We do this by intercepting the component secrets through a subclass
139 that records them.
140 """
141 from i2p_crypto.x25519 import X25519DH
142 from i2p_crypto import mlkem as mlkem_mod
143
144 alice = HybridDHState(MLKEMVariant.ML_KEM_768)
145 alice.generate_keypair()
146
147 bob = HybridDHState(MLKEMVariant.ML_KEM_768)
148 bob.set_remote_public_key(alice.get_public_key())
149
150 response, bob_secret = bob.encapsulate()
151 alice_secret = alice.decapsulate(response)
152
153 # Both sides agree
154 assert alice_secret == bob_secret
155 # Output is 32 bytes (SHA-256)
156 assert len(alice_secret) == 32
157
158 # Verify it's distinct from raw X25519 shared secret alone
159 # (i.e., the ML-KEM contribution matters)
160 alice_x25519_priv = alice._x25519_private
161 bob_x25519_pub = response[:32]
162 raw_x25519_ss = X25519DH.dh(alice_x25519_priv, bob_x25519_pub)
163 assert alice_secret != raw_x25519_ss