distributed.camp
A Python port of the Invisible Internet Project (I2P)
1"""ML-KEM (Module-Lattice Key Encapsulation Mechanism) -- FIPS 203.
2
3Wraps pqcrypto for ML-KEM-512/768/1024. Used in hybrid X25519+ML-KEM
4Noise handshakes for post-quantum forward secrecy.
5"""
6from __future__ import annotations
7
8import enum
9from dataclasses import dataclass
10from typing import Any, Tuple
11
12# Lazy-loaded module references per variant
13_kem_modules: dict[str, Any] = {}
14
15
16def _get_kem_module(module_name: str) -> Any:
17 """Lazy-import a pqcrypto.kem module and cache it."""
18 if module_name not in _kem_modules:
19 try:
20 import importlib
21 _kem_modules[module_name] = importlib.import_module(f"pqcrypto.kem.{module_name}")
22 except ImportError:
23 raise ImportError(
24 "pqcrypto is required for ML-KEM support. "
25 "Install it with: pip install 'i2p-python[pqc]'"
26 )
27 return _kem_modules[module_name]
28
29
30def is_available() -> bool:
31 """Return True if pqcrypto is installed and importable."""
32 try:
33 _get_kem_module("ml_kem_768")
34 return True
35 except ImportError:
36 return False
37
38
39class MLKEMVariant(enum.Enum):
40 """ML-KEM security levels."""
41
42 # pub_len priv_len ct_len ss_len module_name
43 ML_KEM_512 = (800, 1632, 768, 32, "ml_kem_512")
44 ML_KEM_768 = (1184, 2400, 1088, 32, "ml_kem_768")
45 ML_KEM_1024 = (1568, 3168, 1568, 32, "ml_kem_1024")
46
47 def __init__(
48 self,
49 public_key_len: int,
50 private_key_len: int,
51 ciphertext_len: int,
52 shared_secret_len: int,
53 module_name: str,
54 ) -> None:
55 self._public_key_len = public_key_len
56 self._private_key_len = private_key_len
57 self._ciphertext_len = ciphertext_len
58 self._shared_secret_len = shared_secret_len
59 self._module_name = module_name
60
61 @property
62 def public_key_len(self) -> int:
63 return self._public_key_len
64
65 @property
66 def private_key_len(self) -> int:
67 return self._private_key_len
68
69 @property
70 def ciphertext_len(self) -> int:
71 return self._ciphertext_len
72
73 @property
74 def shared_secret_len(self) -> int:
75 return self._shared_secret_len
76
77 @property
78 def module_name(self) -> str:
79 return self._module_name
80
81
82@dataclass(frozen=True)
83class MLKEMKeyPair:
84 """A ML-KEM keypair."""
85 public_key: bytes
86 private_key: bytes
87 variant: MLKEMVariant
88
89
90def generate_keys(variant: MLKEMVariant) -> MLKEMKeyPair:
91 """Generate a new ML-KEM keypair for the given variant."""
92 mod = _get_kem_module(variant.module_name)
93 public_key, private_key = mod.generate_keypair()
94 return MLKEMKeyPair(
95 public_key=bytes(public_key),
96 private_key=bytes(private_key),
97 variant=variant,
98 )
99
100
101def encapsulate(
102 variant: MLKEMVariant, public_key: bytes
103) -> Tuple[bytes, bytes]:
104 """Encapsulate against a public key, returning (ciphertext, shared_secret)."""
105 mod = _get_kem_module(variant.module_name)
106 ciphertext, shared_secret = mod.encrypt(public_key)
107 return bytes(ciphertext), bytes(shared_secret)
108
109
110def decapsulate(
111 variant: MLKEMVariant, ciphertext: bytes, private_key: bytes
112) -> bytes:
113 """Decapsulate a ciphertext with a private key, returning the shared secret."""
114 mod = _get_kem_module(variant.module_name)
115 shared_secret = mod.decrypt(private_key, ciphertext)
116 return bytes(shared_secret)