diff --git a/src/kyber_py/drbg/aes256_ctr_drbg.py b/src/kyber_py/drbg/aes256_ctr_drbg.py index db28e18..5110399 100644 --- a/src/kyber_py/drbg/aes256_ctr_drbg.py +++ b/src/kyber_py/drbg/aes256_ctr_drbg.py @@ -69,7 +69,8 @@ def ctr_drbg_update(self, provided_data): self.V = tmp[32:] def random_bytes(self, num_bytes, additional=None): - if self.reseed_ctr >= self.reseed_interval: + # We don't cover this in coverage as we would need to run the counter 2^48 times + if self.reseed_ctr >= self.reseed_interval: # pragma: no cover raise Warning("The DRBG has been exhausted! Reseed!") # Set the optional additional information diff --git a/src/kyber_py/kyber/kyber.py b/src/kyber_py/kyber/kyber.py index 81b65f4..93b489e 100644 --- a/src/kyber_py/kyber/kyber.py +++ b/src/kyber_py/kyber/kyber.py @@ -44,7 +44,7 @@ def set_drbg_seed(self, seed): self._drbg = AES256_CTR_DRBG(seed) self.random_bytes = self._drbg.random_bytes - except ImportError as e: + except ImportError as e: # pragma: no cover print(f"Error importing AES from pycryptodome: {e = }") raise Warning( "Cannot set DRBG seed due to missing dependencies, try installing requirements: pip -r install requirements" diff --git a/src/kyber_py/ml_kem/ml_kem.py b/src/kyber_py/ml_kem/ml_kem.py index 3708d6a..3eecbce 100644 --- a/src/kyber_py/ml_kem/ml_kem.py +++ b/src/kyber_py/ml_kem/ml_kem.py @@ -49,7 +49,7 @@ def set_drbg_seed(self, seed): self._drbg = AES256_CTR_DRBG(seed) self.random_bytes = self._drbg.random_bytes - except ImportError as e: + except ImportError as e: # pragma: no cover print(f"Error importing AES from pycryptodome: {e = }") raise Warning( "Cannot set DRBG seed due to missing dependencies, try installing requirements: pip -r install requirements" @@ -173,21 +173,29 @@ def _pke_keygen(self): def _pke_encrypt(self, ek_pke, m, r): """ Algorithm 13 + + As well as performing the usual pke encryption, the FIPS document + requires two additional checks. + + 1. Type Check: The ek_pke is of the expected length + 2. Modulus Check: That t_hat has been canonically encoded """ + # These should always hold, as encaps() generates m, r and + # _pke_encrypt should never be called directly by a user assert len(m) == 32 assert len(r) == 32 + # First check if the encap key has the right length + if len(ek_pke) != 384 * self.k + 32: + raise ValueError("Type check failed, ek_pke has the wrong length") + # Unpack ek t_hat_bytes, rho = ek_pke[:-32], ek_pke[-32:] # Compute Polynomial from bytes t_hat = self.M.decode_vector(t_hat_bytes, self.k, 12, is_ntt=True) - # NOTE: - # Perform the input validation checks for ML-KEM - if len(ek_pke) != 384 * self.k + 32: - raise ValueError("Type check failed, ek_pke has the wrong length") - + # 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" diff --git a/tests/test_drbg.py b/tests/test_drbg.py new file mode 100644 index 0000000..3fe3b86 --- /dev/null +++ b/tests/test_drbg.py @@ -0,0 +1,53 @@ +import unittest +import os +from kyber_py.drbg.aes256_ctr_drbg import AES256_CTR_DRBG + + +class TestDRBG(unittest.TestCase): + """ + Some small tests, as the general check for the DRBG is that + the KAT vectors match with the assets within the mlkem and + kyber tests. + """ + + def test_no_seed(self): + # If seed is none, os.urandom is used instead + seed = None + drbg = AES256_CTR_DRBG(seed) + self.assertNotEqual(drbg.entropy_input, None) + + def test_bad_seed(self): + # if the seed length is not 48, the code fails + seed = b"1" + self.assertRaises(ValueError, lambda: AES256_CTR_DRBG(seed)) + seed = b"1" * 49 + self.assertRaises(ValueError, lambda: AES256_CTR_DRBG(seed)) + + def test_personalization(self): + # if the personalization is longer than 48 bytes, fail + seed = os.urandom(48) + personalization = os.urandom(24) + drbg = AES256_CTR_DRBG(seed, personalization) + self.assertEqual(AES256_CTR_DRBG, type(drbg)) + + def test_bad_personalization(self): + # if the personalization is longer than 48 bytes, fail + seed = os.urandom(48) + personalization = os.urandom(49) + self.assertRaises( + ValueError, lambda: AES256_CTR_DRBG(seed, personalization) + ) + + def test_additional(self): + drbg = AES256_CTR_DRBG() + additional = os.urandom(24) + b = drbg.random_bytes(32, additional) + self.assertEqual(len(b), 32) + self.assertEqual(type(b), bytes) + + def test_bad_additional(self): + drbg = AES256_CTR_DRBG() + additional = os.urandom(49) + self.assertRaises( + ValueError, lambda: drbg.random_bytes(32, additional) + ) diff --git a/tests/test_polynomial_generic.py b/tests/test_polynomial_generic.py index 500f4a1..00fd528 100644 --- a/tests/test_polynomial_generic.py +++ b/tests/test_polynomial_generic.py @@ -60,14 +60,16 @@ def test_equality(self): f1 = self.R.random_element() f2 = -f1 self.assertEqual(f1, f1) + # We don't cover the case of f1 being zero, as it's incredibly unlikely to happen if f1.is_zero(): - self.assertTrue(f1 == f2) + self.assertTrue(f1 == f2) # pragma: no cover else: self.assertFalse(f1 == f2) self.assertTrue(self.R(0) == 0) self.assertTrue(self.R(1) == self.R.q + 1) self.assertTrue(self.R(self.R.q - 1) == -1) + self.assertTrue(self.R(0) != 1) self.assertFalse(self.R(self.R.q - 1) == "a") def test_add_failure(self):