attempt to increase coverage

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

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"

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"

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)
+        )

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):