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