add explicit type hints to ml-kem

Author: GiacomoPope

benchmarks/benchmark_kyber.py

@@ -53,9 +53,9 @@ def benchmark_kyber(Kyber, name, count):
     avg_dec = sum(dec_times) / count
     print(
         f" {name:11} |"
-        f"{avg_keygen*1000:7.2f}ms | {1/avg_keygen:10.2f} |"
-        f"{avg_enc*1000:6.2f}ms | {1/avg_enc:9.2f} |"
-        f"{avg_dec*1000:6.2f}ms | {1/avg_dec:7.2f} |"
+        f"{avg_keygen * 1000:7.2f}ms | {1 / avg_keygen:10.2f} |"
+        f"{avg_enc * 1000:6.2f}ms | {1 / avg_enc:9.2f} |"
+        f"{avg_dec * 1000:6.2f}ms | {1 / avg_dec:7.2f} |"
     )
 
 

benchmarks/benchmark_ml_kem.py

@@ -53,9 +53,9 @@ def benchmark_ml_kem(ML_KEM, name, count):
     avg_dec = sum(dec_times) / count
     print(
         f" {name:11} |"
-        f"{avg_keygen*1000:7.2f}ms | {1/avg_keygen:10.2f} |"
-        f"{avg_enc*1000:6.2f}ms | {1/avg_enc:9.2f} |"
-        f"{avg_dec*1000:6.2f}ms | {1/avg_dec:7.2f} |"
+        f"{avg_keygen * 1000:7.2f}ms | {1 / avg_keygen:10.2f} |"
+        f"{avg_enc * 1000:6.2f}ms | {1 / avg_enc:9.2f} |"
+        f"{avg_dec * 1000:6.2f}ms | {1 / avg_dec:7.2f} |"
     )
 
 

src/kyber_py/drbg/aes256_ctr_drbg.py

@@ -5,9 +5,7 @@
 
 
 class AES256_CTR_DRBG:
-    def __init__(
-        self, seed: Optional[bytes] = None, personalization: bytes = b""
-    ):
+    def __init__(self, seed: Optional[bytes] = None, personalization: bytes = b""):
         """
         DRBG implementation based on AES-256 CTR following the document NIST SP
         800-90A Section 10.2.1
@@ -59,9 +57,7 @@ def __instantiate(self, personalization: bytes = b"") -> bytes:
                 f"{self.seed_length}. Input has length {len(personalization)}"
             )
         # Ensure personalization has exactly seed_length bytes
-        personalization += bytes([0]) * (
-            self.seed_length - len(personalization)
-        )
+        personalization += bytes([0]) * (self.seed_length - len(personalization))
         # debugging
         assert len(personalization) == self.seed_length
         return xor_bytes(self.entropy_input, personalization)
@@ -97,9 +93,7 @@ def __ctr_drbg_update(self, provided_data: bytes) -> None:
         self.key = tmp[:32]
         self.V = tmp[32:]
 
-    def random_bytes(
-        self, num_bytes: int, additional: Optional[bytes] = None
-    ) -> bytes:
+    def random_bytes(self, num_bytes: int, additional: Optional[bytes] = None) -> bytes:
         """
         Generate pseudorandom bytes without a generating function
 

src/kyber_py/kyber/kyber.py

@@ -1,6 +1,6 @@
 import os
 from hashlib import sha3_256, sha3_512, shake_128, shake_256
-from ..modules.modules import ModuleKyber
+from ..modules.modules import Module
 from ..utilities.utils import select_bytes
 
 
@@ -17,7 +17,7 @@ def __init__(self, parameter_set):
         self.du = parameter_set["du"]
         self.dv = parameter_set["dv"]
 
-        self.M = ModuleKyber()
+        self.M = Module()
         self.R = self.M.ring
 
         # Use system randomness by default, for deterministic randomness
@@ -111,7 +111,7 @@ def _generate_error_vector(self, sigma, eta, N):
         Helper function which generates a element in the
         module from the Centered Binomial Distribution.
         """
-        elements = [0 for _ in range(self.k)]
+        elements = [self.R.zero() for _ in range(self.k)]
         for i in range(self.k):
             input_bytes = self._prf(sigma, bytes([N]), 64 * eta)
             elements[i] = self.R.cbd(input_bytes, eta)
@@ -135,7 +135,7 @@ def _generate_matrix_from_seed(self, rho, transpose=False):
 
         When `transpose` is set to True, the matrix A is built as the transpose.
         """
-        A_data = [[0 for _ in range(self.k)] for _ in range(self.k)]
+        A_data = [[self.R.zero() for _ in range(self.k)] for _ in range(self.k)]
         for i in range(self.k):
             for j in range(self.k):
                 input_bytes = self._xof(rho, bytes([j]), bytes([i]))

src/kyber_py/ml_kem/ml_kem.py

@@ -5,12 +5,13 @@
 
 import os
 from hashlib import sha3_256, sha3_512, shake_128, shake_256
-from ..modules.modules import ModuleKyber
+from ..modules.modules import Module, Matrix, Vector
+from ..polynomials.polynomials import Polynomial
 from ..utilities.utils import select_bytes
 
 
 class ML_KEM:
-    def __init__(self, params):
+    def __init__(self, params: dict):
         """
         Initialise the ML-KEM with specified lattice parameters.
 
@@ -23,31 +24,31 @@ def __init__(self, params):
         self.du = params["du"]
         self.dv = params["dv"]
 
-        self.M = ModuleKyber()
+        self.M = Module()
         self.R = self.M.ring
         self.oid = params["oid"] if "oid" in params else None
 
         # Use system randomness by default, for deterministic randomness
         # use the method `set_drbg_seed()`
         self.random_bytes = os.urandom
 
-    def _ek_size(self):
+    def _ek_size(self) -> int:
         """
         Return the size of the encapsulation key for the selected paramters.
 
         :rtype: int
         """
         return 384 * self.k + 32
 
-    def _dk_size(self):
+    def _dk_size(self) -> int:
         """
         Return the size of the decapsulation key for the selected parameters.
 
         :rtype: int
         """
         return 768 * self.k + 96
 
-    def set_drbg_seed(self, seed):
+    def set_drbg_seed(self, seed: bytes):
         """
         Change entropy source to a DRBG and seed it with provided value.
 
@@ -74,7 +75,7 @@ def set_drbg_seed(self, seed):
             )
 
     @staticmethod
-    def _xof(bytes32, i, j):
+    def _xof(b: bytes, i: bytes, j: bytes) -> bytes:
         """
         eXtendable-Output Function (XOF) described in 4.9 of FIPS 203 (page 19)
 
@@ -88,15 +89,15 @@ def _xof(bytes32, i, j):
           Casa de Chá da Boa Nova
           https://cryptojedi.org/papers/terminate-20230516.pdf
         """
-        input_bytes = bytes32 + i + j
+        input_bytes = b + i + j
         if len(input_bytes) != 34:
             raise ValueError(
                 "Input bytes should be one 32 byte array and 2 single bytes."
             )
         return shake_128(input_bytes).digest(840)
 
     @staticmethod
-    def _prf(eta, s, b):
+    def _prf(eta: int, s: bytes, b: bytes) -> bytes:
         """
         Pseudorandom function described in 4.3 of FIPS 203 (page 18)
         """
@@ -108,57 +109,61 @@ def _prf(eta, s, b):
         return shake_256(input_bytes).digest(eta * 64)
 
     @staticmethod
-    def _H(s):
+    def _H(s: bytes) -> bytes:
         """
         Hash function described in 4.4 of FIPS 203 (page 18)
         """
         return sha3_256(s).digest()
 
     @staticmethod
-    def _J(s):
+    def _J(s: bytes) -> bytes:
         """
         Hash function described in 4.4 of FIPS 203 (page 18)
         """
         return shake_256(s).digest(32)
 
     @staticmethod
-    def _G(s):
+    def _G(s: bytes) -> tuple[bytes, bytes]:
         """
         Hash function described in 4.5 of FIPS 203 (page 18)
         """
         h = sha3_512(s).digest()
         return h[:32], h[32:]
 
-    def _generate_matrix_from_seed(self, rho, transpose=False):
+    def _generate_matrix_from_seed(self, rho: bytes, transpose: bool = False) -> Matrix:
         """
         Helper function which generates a element of size
         k x k from a seed `rho`.
 
         When `transpose` is set to True, the matrix A is
         built as the transpose.
         """
-        A_data = [[0 for _ in range(self.k)] for _ in range(self.k)]
+        A_data = [[self.R.zero() for _ in range(self.k)] for _ in range(self.k)]
         for i in range(self.k):
             for j in range(self.k):
                 xof_bytes = self._xof(rho, bytes([j]), bytes([i]))
                 A_data[i][j] = self.R.ntt_sample(xof_bytes)
         A_hat = self.M(A_data, transpose=transpose)
         return A_hat
 
-    def _generate_error_vector(self, sigma, eta, N):
+    def _generate_error_vector(
+        self, sigma: bytes, eta: int, N: int
+    ) -> tuple[Vector, int]:
         """
         Helper function which generates a element in the
         module from the Centered Binomial Distribution.
         """
-        elements = [0 for _ in range(self.k)]
+        elements = [self.R.zero() for _ in range(self.k)]
         for i in range(self.k):
             prf_output = self._prf(eta, sigma, bytes([N]))
             elements[i] = self.R.cbd(prf_output, eta)
             N += 1
         v = self.M.vector(elements)
         return v, N
 
-    def _generate_polynomial(self, sigma, eta, N):
+    def _generate_polynomial(
+        self, sigma: bytes, eta: int, N: int
+    ) -> tuple[Polynomial, int]:
         """
         Helper function which generates a element in the
         polynomial ring from the Centered Binomial Distribution.
@@ -167,7 +172,7 @@ def _generate_polynomial(self, sigma, eta, N):
         p = self.R.cbd(prf_output, eta)
         return p, N + 1
 
-    def _k_pke_keygen(self, d):
+    def _k_pke_keygen(self, d: bytes) -> tuple[bytes, bytes]:
         """
         Use randomness to generate an encryption key and a corresponding
         decryption key following Algorithm 13 (FIPS 203)
@@ -203,7 +208,7 @@ def _k_pke_keygen(self, d):
 
         return (ek_pke, dk_pke)
 
-    def _k_pke_encrypt(self, ek_pke, m, r):
+    def _k_pke_encrypt(self, ek_pke: bytes, m: bytes, r: bytes) -> bytes:
         """
         Uses the encryption key to encrypt a plaintext message using the
         randomness r following Algorithm 14 (FIPS 203)
@@ -231,9 +236,7 @@ def _k_pke_encrypt(self, ek_pke, m, r):
 
         # Next check that t_hat has been canonically encoded
         if t_hat.encode(12) != t_hat_bytes:
-            raise ValueError(
-                "Modulus check failed, t_hat does not encode correctly"
-            )
+            raise ValueError("Modulus check failed, t_hat does not encode correctly")
 
         # Generate A_hat^T from seed rho
         A_hat_T = self._generate_matrix_from_seed(rho, transpose=True)
@@ -255,7 +258,7 @@ def _k_pke_encrypt(self, ek_pke, m, r):
 
         return c1 + c2
 
-    def _k_pke_decrypt(self, dk_pke, c):
+    def _k_pke_decrypt(self, dk_pke: bytes, c: bytes) -> bytes:
         """
         Uses the decryption key to decrypt a ciphertext following
         Algorithm 15 (FIPS 203)
@@ -273,7 +276,7 @@ def _k_pke_decrypt(self, dk_pke, c):
 
         return m
 
-    def _keygen_internal(self, d, z):
+    def _keygen_internal(self, d: bytes, z: bytes) -> tuple[bytes, bytes]:
         """
         Use randomness to generate an encapsulation key and a corresponding
         decapsulation key following Algorithm 16 (FIPS 203)
@@ -288,7 +291,7 @@ def _keygen_internal(self, d, z):
 
         return (ek, dk)
 
-    def keygen(self):
+    def keygen(self) -> tuple[bytes, bytes]:
         """
         Generate an encapsulation key and corresponding decapsulation key
         following Algorithm 19 (FIPS 203)
@@ -309,7 +312,7 @@ def keygen(self):
         ) = self._keygen_internal(d, z)
         return (ek, dk)
 
-    def key_derive(self, seed):
+    def key_derive(self, seed: bytes) -> tuple[bytes, bytes]:
         """
         Derive an encapsulation key and corresponding decapsulation key
         following the approach from Section 7.1 (FIPS 203)
@@ -329,7 +332,7 @@ def key_derive(self, seed):
         ek, dk = self._keygen_internal(d, z)
         return (ek, dk)
 
-    def _encaps_internal(self, ek, m):
+    def _encaps_internal(self, ek: bytes, m: bytes) -> tuple[bytes, bytes]:
         """
         Uses the encapsulation key and randomness to generate a key and an
         associated ciphertext following Algorithm 17 (FIPS 203)
@@ -355,7 +358,7 @@ def _encaps_internal(self, ek, m):
 
         return K, c
 
-    def encaps(self, ek):
+    def encaps(self, ek: bytes) -> tuple[bytes, bytes]:
         """
         Uses the encapsulation key to generate a shared secret key and an
         associated ciphertext following Algorithm 20 (FIPS 203)
@@ -374,7 +377,7 @@ def encaps(self, ek):
         K, c = self._encaps_internal(ek, m)
         return K, c
 
-    def _decaps_internal(self, dk, c):
+    def _decaps_internal(self, dk: bytes, c: bytes) -> bytes:
         """
         Uses the decapsulation key to produce a shared secret key from a
         ciphertext following Algorithm 18 (FIPS 203)
@@ -429,7 +432,7 @@ def _decaps_internal(self, dk, c):
         # performed in constant time
         return select_bytes(K_bar, K_prime, c == c_prime)
 
-    def decaps(self, dk, c):
+    def decaps(self, dk: bytes, c: bytes) -> bytes:
         """
         Uses the decapsulation key to produce a shared secret key from a
         ciphertext following Algorithm 21 (FIPS 203).

src/kyber_py/ml_kem/pkcs.py

@@ -34,7 +34,7 @@ def ek_to_der(kem, ek):
         raise ValueError("Only KEMs with specified OIDs can be encoded")
 
     if len(ek) != kem._ek_size():
-        raise ValueError(f"Provided key size doesn't match the provided kem")
+        raise ValueError("Provided key size doesn't match the provided kem")
 
     enc = der.encode_sequence(
         der.encode_sequence(
@@ -252,7 +252,7 @@ def dk_from_der(enc_key):
     else:
         tag, seed, empty = der.remove_implicit(priv_key)
         if tag != 0:
-            raise der.UnexpectedDER(f"Unexpected tag in private key encoding")
+            raise der.UnexpectedDER("Unexpected tag in private key encoding")
         if empty:
             raise der.UnexpectedDER("Junk after seed encoding")