A Python port of the Invisible Internet Project (I2P)
1"""Tests for ML-KEM (FIPS 203) wrapper around pqcrypto.
2
3TDD: tests written before implementation.
4"""
5from __future__ import annotations
6
7import pytest
8
9from i2p_crypto.mlkem import (
10 MLKEMVariant,
11 MLKEMKeyPair,
12 generate_keys,
13 encapsulate,
14 decapsulate,
15 is_available,
16)
17
18# Skip all tests if pqcrypto is not installed
19pytestmark = pytest.mark.skipif(
20 not is_available(), reason="pqcrypto not installed"
21)
22
23
24# ---------- key generation ----------
25
26class TestGenerateKeys:
27 def test_generate_keys_512(self) -> None:
28 kp = generate_keys(MLKEMVariant.ML_KEM_512)
29 assert isinstance(kp, MLKEMKeyPair)
30 assert kp.variant is MLKEMVariant.ML_KEM_512
31 assert len(kp.public_key) == MLKEMVariant.ML_KEM_512.public_key_len
32 assert len(kp.private_key) == MLKEMVariant.ML_KEM_512.private_key_len
33
34 def test_generate_keys_768(self) -> None:
35 kp = generate_keys(MLKEMVariant.ML_KEM_768)
36 assert isinstance(kp, MLKEMKeyPair)
37 assert kp.variant is MLKEMVariant.ML_KEM_768
38 assert len(kp.public_key) == MLKEMVariant.ML_KEM_768.public_key_len
39 assert len(kp.private_key) == MLKEMVariant.ML_KEM_768.private_key_len
40
41 def test_generate_keys_1024(self) -> None:
42 kp = generate_keys(MLKEMVariant.ML_KEM_1024)
43 assert isinstance(kp, MLKEMKeyPair)
44 assert kp.variant is MLKEMVariant.ML_KEM_1024
45 assert len(kp.public_key) == MLKEMVariant.ML_KEM_1024.public_key_len
46 assert len(kp.private_key) == MLKEMVariant.ML_KEM_1024.private_key_len
47
48
49# ---------- encaps / decaps roundtrip ----------
50
51class TestEncapsDecaps:
52 @pytest.mark.parametrize("variant", list(MLKEMVariant))
53 def test_roundtrip(self, variant: MLKEMVariant) -> None:
54 kp = generate_keys(variant)
55 ciphertext, shared_secret_enc = encapsulate(variant, kp.public_key)
56 shared_secret_dec = decapsulate(variant, ciphertext, kp.private_key)
57
58 assert len(ciphertext) == variant.ciphertext_len
59 assert len(shared_secret_enc) == variant.shared_secret_len
60 assert len(shared_secret_dec) == variant.shared_secret_len
61 assert shared_secret_enc == shared_secret_dec
62
63 def test_encaps_decaps_roundtrip_512(self) -> None:
64 kp = generate_keys(MLKEMVariant.ML_KEM_512)
65 ct, ss_enc = encapsulate(MLKEMVariant.ML_KEM_512, kp.public_key)
66 ss_dec = decapsulate(MLKEMVariant.ML_KEM_512, ct, kp.private_key)
67 assert ss_enc == ss_dec
68
69 def test_encaps_decaps_roundtrip_768(self) -> None:
70 kp = generate_keys(MLKEMVariant.ML_KEM_768)
71 ct, ss_enc = encapsulate(MLKEMVariant.ML_KEM_768, kp.public_key)
72 ss_dec = decapsulate(MLKEMVariant.ML_KEM_768, ct, kp.private_key)
73 assert ss_enc == ss_dec
74
75 def test_encaps_decaps_roundtrip_1024(self) -> None:
76 kp = generate_keys(MLKEMVariant.ML_KEM_1024)
77 ct, ss_enc = encapsulate(MLKEMVariant.ML_KEM_1024, kp.public_key)
78 ss_dec = decapsulate(MLKEMVariant.ML_KEM_1024, ct, kp.private_key)
79 assert ss_enc == ss_dec
80
81
82# ---------- security properties ----------
83
84class TestSecurityProperties:
85 def test_different_keypairs_different_secrets(self) -> None:
86 """Two encapsulations with different keys must produce different shared secrets."""
87 kp1 = generate_keys(MLKEMVariant.ML_KEM_768)
88 kp2 = generate_keys(MLKEMVariant.ML_KEM_768)
89
90 _, ss1 = encapsulate(MLKEMVariant.ML_KEM_768, kp1.public_key)
91 _, ss2 = encapsulate(MLKEMVariant.ML_KEM_768, kp2.public_key)
92
93 # With overwhelming probability these differ
94 assert ss1 != ss2
95
96 def test_wrong_private_key_fails(self) -> None:
97 """ML-KEM uses implicit rejection: wrong key yields a different (random) shared secret."""
98 kp1 = generate_keys(MLKEMVariant.ML_KEM_768)
99 kp2 = generate_keys(MLKEMVariant.ML_KEM_768)
100
101 ct, ss_enc = encapsulate(MLKEMVariant.ML_KEM_768, kp1.public_key)
102 ss_wrong = decapsulate(MLKEMVariant.ML_KEM_768, ct, kp2.private_key)
103
104 # Implicit rejection: decapsulate succeeds but returns a different secret
105 assert len(ss_wrong) == MLKEMVariant.ML_KEM_768.shared_secret_len
106 assert ss_wrong != ss_enc
107
108
109# ---------- variant properties ----------
110
111class TestVariantProperties:
112 def test_key_sizes_512(self) -> None:
113 assert MLKEMVariant.ML_KEM_512.public_key_len == 800
114 assert MLKEMVariant.ML_KEM_512.private_key_len == 1632
115 assert MLKEMVariant.ML_KEM_512.ciphertext_len == 768
116 assert MLKEMVariant.ML_KEM_512.shared_secret_len == 32
117
118 def test_key_sizes_768(self) -> None:
119 assert MLKEMVariant.ML_KEM_768.public_key_len == 1184
120 assert MLKEMVariant.ML_KEM_768.private_key_len == 2400
121 assert MLKEMVariant.ML_KEM_768.ciphertext_len == 1088
122 assert MLKEMVariant.ML_KEM_768.shared_secret_len == 32
123
124 def test_key_sizes_1024(self) -> None:
125 assert MLKEMVariant.ML_KEM_1024.public_key_len == 1568
126 assert MLKEMVariant.ML_KEM_1024.private_key_len == 3168
127 assert MLKEMVariant.ML_KEM_1024.ciphertext_len == 1568
128 assert MLKEMVariant.ML_KEM_1024.shared_secret_len == 32
129
130 @pytest.mark.parametrize("variant", list(MLKEMVariant))
131 def test_generated_key_sizes_match_properties(self, variant: MLKEMVariant) -> None:
132 """Verify that actual generated key sizes match the enum properties."""
133 kp = generate_keys(variant)
134 assert len(kp.public_key) == variant.public_key_len
135 assert len(kp.private_key) == variant.private_key_len
136
137 def test_module_names(self) -> None:
138 assert MLKEMVariant.ML_KEM_512.module_name == "ml_kem_512"
139 assert MLKEMVariant.ML_KEM_768.module_name == "ml_kem_768"
140 assert MLKEMVariant.ML_KEM_1024.module_name == "ml_kem_1024"
141
142
143# ---------- availability ----------
144
145class TestAvailability:
146 @pytest.mark.skipif(False, reason="always runs")
147 def test_is_available(self) -> None:
148 """is_available() returns True when pqcrypto is installed (this test only runs when it is)."""
149 assert is_available() is True