use secrets module instead of random for better security

PyryL · Dec 12, 2023 · f1bcc0e · f1bcc0e
1 parent ebd0837
commit f1bcc0e
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 8 deletions.
diff --git a/kyber/ccakem.py b/kyber/ccakem.py
@@ -1,4 +1,4 @@
-from random import randbytes
+from secrets import token_bytes
 from kyber.encryption import generate_keys, Encrypt, Decrypt
 from kyber.utils.pseudo_random import H, G, kdf
 from kyber.constants import k, n, du, dv
@@ -9,7 +9,7 @@ def ccakem_generate_keys() -> tuple[bytes, bytes]:
     :returns (private_key, public_key) tuple
     """
 
-    z = randbytes(32)
+    z = token_bytes(32)
     sk, pk = generate_keys()
     sk = sk + pk + H(pk) + z
 
@@ -29,7 +29,7 @@ def ccakem_encrypt(public_key: bytes, shared_secret_length: int = 32) -> tuple[b
 
     assert len(public_key) == 12 * k * n//8 + 32
 
-    m = H(randbytes(32))
+    m = H(token_bytes(32))
     Kr = G(m + H(public_key))
     K, r = Kr[:32], Kr[32:]
     c = Encrypt(public_key, m, r).encrypt()

diff --git a/kyber/encryption/encrypt.py b/kyber/encryption/encrypt.py
@@ -1,4 +1,4 @@
-from random import randbytes
+from secrets import token_bytes
 import numpy as np
 from kyber.utils.cbd import cbd
 from kyber.utils.pseudo_random import prf
@@ -13,8 +13,8 @@
 class Encrypt:
     def __init__(self, public_key: bytes, m: bytes = None, r: bytes = None) -> None:
         self._pk = public_key
-        self._m = m if m is not None else randbytes(32)
-        self._r = r if r is not None else randbytes(32)
+        self._m = m if m is not None else token_bytes(32)
+        self._r = r if r is not None else token_bytes(32)
         assert len(self._m) == 32
         assert len(self._r) == 32
         if len(self._pk) != 12 * k * int(n/8) + 32:

diff --git a/kyber/encryption/keygen.py b/kyber/encryption/keygen.py
@@ -1,4 +1,4 @@
-from random import randbytes
+from secrets import token_bytes
 import numpy as np
 from kyber.constants import k, eta1
 from kyber.utils.pseudo_random import prf, G, xof
@@ -14,7 +14,7 @@ def generate_keys() -> tuple:
     :returns (private_key, public_key)
     """
 
-    d = randbytes(32)
+    d = token_bytes(32)
     rho, sigma = G(d)[:32], G(d)[32:]
 
     A = np.empty((k, k), PolynomialRing)