A Python port of the Invisible Internet Project (I2P)
at main 163 lines 5.6 kB view raw
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