From 305d1b653ee04e49ea38abf549fd45436ba0b799 Mon Sep 17 00:00:00 2001 From: Giacomo Pope Date: Sat, 20 Jul 2024 22:22:47 +0100 Subject: [PATCH] restructure code --- __init__.py | 0 .../benchmark_kyber.py | 0 benchmarks/benchmark_ml_kem.py | 73 +++++++++++++++++++ docs/source/api.rst | 8 +- docs/source/{aes256_ctr_drbg.rst => drbg.rst} | 4 +- docs/source/modules_generic.rst | 7 -- docs/source/polynomials_generic.rst | 7 -- docs/source/run_kyber.rst | 7 -- docs/source/{utils.rst => utilities.rst} | 4 +- drbg/__init__.py | 0 aes256_ctr_drbg.py => drbg/aes256_ctr_drbg.py | 4 +- kyber.py | 7 +- ml_kem.py | 7 +- modules/__init__.py | 0 modules.py => modules/modules.py | 4 +- .../modules_generic.py | 0 polynomials/__init__.py | 0 polynomials.py => polynomials/polynomials.py | 4 +- .../polynomials_generic.py | 0 run_kyber.py | 8 -- test_ml_kem.py | 28 ------- tests/__init__.py | 0 test_kyber.py => tests/test_kyber.py | 2 +- test_kyber_kat.py => tests/test_kyber_kat.py | 8 +- tests/test_ml_kem.py | 30 ++++++++ utilities/__init__.py | 0 utils.py => utilities/utils.py | 0 27 files changed, 129 insertions(+), 83 deletions(-) create mode 100644 __init__.py rename benchmark_kyber.py => benchmarks/benchmark_kyber.py (100%) create mode 100644 benchmarks/benchmark_ml_kem.py rename docs/source/{aes256_ctr_drbg.rst => drbg.rst} (58%) delete mode 100644 docs/source/modules_generic.rst delete mode 100644 docs/source/polynomials_generic.rst delete mode 100644 docs/source/run_kyber.rst rename docs/source/{utils.rst => utilities.rst} (61%) create mode 100644 drbg/__init__.py rename aes256_ctr_drbg.py => drbg/aes256_ctr_drbg.py (98%) create mode 100644 modules/__init__.py rename modules.py => modules/modules.py (95%) rename modules_generic.py => modules/modules_generic.py (100%) create mode 100644 polynomials/__init__.py rename polynomials.py => polynomials/polynomials.py (98%) rename polynomials_generic.py => polynomials/polynomials_generic.py (100%) delete mode 100644 run_kyber.py delete mode 100644 test_ml_kem.py create mode 100644 tests/__init__.py rename test_kyber.py => tests/test_kyber.py (99%) rename test_kyber_kat.py => tests/test_kyber_kat.py (90%) create mode 100644 tests/test_ml_kem.py create mode 100644 utilities/__init__.py rename utils.py => utilities/utils.py (100%) diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/benchmark_kyber.py b/benchmarks/benchmark_kyber.py similarity index 100% rename from benchmark_kyber.py rename to benchmarks/benchmark_kyber.py diff --git a/benchmarks/benchmark_ml_kem.py b/benchmarks/benchmark_ml_kem.py new file mode 100644 index 0000000..85d8631 --- /dev/null +++ b/benchmarks/benchmark_ml_kem.py @@ -0,0 +1,73 @@ +from ml_kem import ML_KEM128, ML_KEM192, ML_KEM256 +import cProfile +from time import time + + +def profile_ml_kem(ML_KEM): + (ek, dk) = ML_KEM.keygen() + (K, c) = ML_KEM.encaps(ek) + + gvars = {} + lvars = {"ML_KEM": ML_KEM, "c": c, "ek": ek, "dk": dk} + + cProfile.runctx( + "[ML_KEM.keygen() for _ in range(100)]", + globals=gvars, + locals=lvars, + sort=1, + ) + cProfile.runctx( + "[ML_KEM.encaps(ek) for _ in range(100)]", + globals=gvars, + locals=lvars, + sort=1, + ) + cProfile.runctx( + "[ML_KEM.decaps(c, dk) for _ in range(100)]", + globals=gvars, + locals=lvars, + sort=1, + ) + + +def benchmark_ml_kem(ML_KEM, name, count): + keygen_times = [] + enc_times = [] + dec_times = [] + + for _ in range(count): + t0 = time() + ek, dk = ML_KEM.keygen() + keygen_times.append(time() - t0) + + t1 = time() + _, c = ML_KEM.encaps(ek) + enc_times.append(time() - t1) + + t2 = time() + _ = ML_KEM.decaps(c, dk) + dec_times.append(time() - t2) + + avg_keygen = sum(keygen_times) / count + avg_enc = sum(enc_times) / count + avg_dec = sum(dec_times) / count + print( + f" {name:11} |" + f"{avg_keygen*1000:8.2f}ms {1/avg_keygen:11.2f}" + f"{avg_enc*1000:8.2f}ms {1/avg_enc:10.2f}" + f"{avg_dec*1000:8.2f}ms {1/avg_dec:8.2f}" + ) + + +if __name__ == "__main__": + count = 1000 + # common banner + print("-" * 80) + print( + " Params | keygen | keygen/s | encap | encap/s " + "| decap | decap/s" + ) + print("-" * 80) + benchmark_ml_kem(ML_KEM128, "ML_KEM128", count) + benchmark_ml_kem(ML_KEM192, "ML_KEM192", count) + benchmark_ml_kem(ML_KEM256, "ML_KEM256", count) diff --git a/docs/source/api.rst b/docs/source/api.rst index 7606aea..93806ff 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -2,13 +2,9 @@ API === .. toctree:: - aes256_ctr_drbg - benchmark_kyber + drbg kyber ml_kem modules - modules_generic polynomials - polynomials_generic - run_kyber - utils + utilities diff --git a/docs/source/aes256_ctr_drbg.rst b/docs/source/drbg.rst similarity index 58% rename from docs/source/aes256_ctr_drbg.rst rename to docs/source/drbg.rst index 4dcf0f0..d3a8b8d 100644 --- a/docs/source/aes256_ctr_drbg.rst +++ b/docs/source/drbg.rst @@ -1,7 +1,7 @@ -aes256\_ctr\_drbg module +drbg module ======================== -.. automodule:: aes256_ctr_drbg +.. automodule:: drbg :members: :undoc-members: :show-inheritance: diff --git a/docs/source/modules_generic.rst b/docs/source/modules_generic.rst deleted file mode 100644 index ce18b16..0000000 --- a/docs/source/modules_generic.rst +++ /dev/null @@ -1,7 +0,0 @@ -modules\_generic module -======================= - -.. automodule:: modules_generic - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/source/polynomials_generic.rst b/docs/source/polynomials_generic.rst deleted file mode 100644 index 0be0ea3..0000000 --- a/docs/source/polynomials_generic.rst +++ /dev/null @@ -1,7 +0,0 @@ -polynomials\_generic module -=========================== - -.. automodule:: polynomials_generic - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/source/run_kyber.rst b/docs/source/run_kyber.rst deleted file mode 100644 index c1ab145..0000000 --- a/docs/source/run_kyber.rst +++ /dev/null @@ -1,7 +0,0 @@ -run\_kyber module -================= - -.. automodule:: run_kyber - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/source/utils.rst b/docs/source/utilities.rst similarity index 61% rename from docs/source/utils.rst rename to docs/source/utilities.rst index 44cef9e..4cda6fa 100644 --- a/docs/source/utils.rst +++ b/docs/source/utilities.rst @@ -1,7 +1,7 @@ -utils module +utilities module ============ -.. automodule:: utils +.. automodule:: utilities :members: :undoc-members: :show-inheritance: diff --git a/drbg/__init__.py b/drbg/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/aes256_ctr_drbg.py b/drbg/aes256_ctr_drbg.py similarity index 98% rename from aes256_ctr_drbg.py rename to drbg/aes256_ctr_drbg.py index 0cc9391..7ae6c2f 100644 --- a/aes256_ctr_drbg.py +++ b/drbg/aes256_ctr_drbg.py @@ -1,5 +1,5 @@ import os -from utils import xor_bytes +from utilities.utils import xor_bytes from Crypto.Cipher import AES @@ -88,7 +88,7 @@ def random_bytes(self, num_bytes, additional=None): if len(additional) > self.seed_length: raise ValueError( f"The additional input must be of length at most: " - f"{self.seed_length}. Input has length {len(seed)}" + f"{self.seed_length}. Input has length {len(additional)}" ) elif len(additional) < self.seed_length: additional += bytes([0]) * (self.seed_length - len(additional)) diff --git a/kyber.py b/kyber.py index c2dd03d..ed088b4 100644 --- a/kyber.py +++ b/kyber.py @@ -1,10 +1,9 @@ import os from hashlib import sha3_256, sha3_512, shake_128, shake_256 -from polynomials import PolynomialRingKyber -from modules import ModuleKyber +from modules.modules import ModuleKyber try: - from aes256_ctr_drbg import AES256_CTR_DRBG + from drbg.aes256_ctr_drbg import AES256_CTR_DRBG except ImportError as e: print( "Error importing AES CTR DRBG. Have you tried installing requirements?" @@ -47,8 +46,8 @@ def __init__(self, parameter_set): self.du = parameter_set["du"] self.dv = parameter_set["dv"] - self.R = PolynomialRingKyber() self.M = ModuleKyber() + self.R = self.M.ring self.drbg = None self.random_bytes = os.urandom diff --git a/ml_kem.py b/ml_kem.py index 73013de..730e4a2 100644 --- a/ml_kem.py +++ b/ml_kem.py @@ -1,10 +1,9 @@ import os from hashlib import sha3_256, sha3_512, shake_128, shake_256 -from polynomials import PolynomialRingKyber -from modules import ModuleKyber +from modules.modules import ModuleKyber try: - from aes256_ctr_drbg import AES256_CTR_DRBG + from drbg.aes256_ctr_drbg import AES256_CTR_DRBG except ImportError as e: print( "Error importing AES CTR DRBG. Have you tried installing requirements?" @@ -32,8 +31,8 @@ def __init__(self, params, seed=None): self.du = params["du"] self.dv = params["dv"] - self.R = PolynomialRingKyber() self.M = ModuleKyber() + self.R = self.M.ring # NIST approved randomness if seed is None: diff --git a/modules/__init__.py b/modules/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/modules.py b/modules/modules.py similarity index 95% rename from modules.py rename to modules/modules.py index ccafb81..6ec9722 100644 --- a/modules.py +++ b/modules/modules.py @@ -1,5 +1,5 @@ -from polynomials import PolynomialRingKyber -from modules_generic import Module, Matrix +from polynomials.polynomials import PolynomialRingKyber +from modules.modules_generic import Module, Matrix class ModuleKyber(Module): diff --git a/modules_generic.py b/modules/modules_generic.py similarity index 100% rename from modules_generic.py rename to modules/modules_generic.py diff --git a/polynomials/__init__.py b/polynomials/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/polynomials.py b/polynomials/polynomials.py similarity index 98% rename from polynomials.py rename to polynomials/polynomials.py index cb20712..9437e8b 100644 --- a/polynomials.py +++ b/polynomials/polynomials.py @@ -1,5 +1,5 @@ -from polynomials_generic import PolynomialRing, Polynomial -from utils import bytes_to_bits, bitstring_to_bytes +from polynomials.polynomials_generic import PolynomialRing, Polynomial +from utilities.utils import bytes_to_bits, bitstring_to_bytes class PolynomialRingKyber(PolynomialRing): diff --git a/polynomials_generic.py b/polynomials/polynomials_generic.py similarity index 100% rename from polynomials_generic.py rename to polynomials/polynomials_generic.py diff --git a/run_kyber.py b/run_kyber.py deleted file mode 100644 index a4cc2d0..0000000 --- a/run_kyber.py +++ /dev/null @@ -1,8 +0,0 @@ -from kyber import Kyber512, Kyber768, Kyber1024 - -for _ in range(10): - pk, sk = Kyber512.keygen() - for _ in range(10): - c, key = Kyber512.enc(pk) - _key = Kyber512.dec(c, sk) - assert key == _key diff --git a/test_ml_kem.py b/test_ml_kem.py deleted file mode 100644 index 71a6640..0000000 --- a/test_ml_kem.py +++ /dev/null @@ -1,28 +0,0 @@ -from ml_kem import ML_KEM128, ML_KEM192, ML_KEM256 - - -def _test_ml_kem(ml_kem): - for _ in range(50): - (ek, dk) = ml_kem.keygen() - (K, c) = ml_kem.encaps(ek) - K_prime = ml_kem.decaps(c, dk) - - assert K == K_prime - - -def test_ml_kem_pke_128(): - _test_ml_kem(ML_KEM128) - - -def test_ml_kem_pke_192(): - _test_ml_kem(ML_KEM192) - - -def test_ml_kem_pke_256(): - _test_ml_kem(ML_KEM256) - - -if __name__ == "__main__": - test_ml_kem_pke_128() - test_ml_kem_pke_192() - test_ml_kem_pke_256() diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/test_kyber.py b/tests/test_kyber.py similarity index 99% rename from test_kyber.py rename to tests/test_kyber.py index ccd2264..604554a 100644 --- a/test_kyber.py +++ b/tests/test_kyber.py @@ -1,7 +1,7 @@ import unittest import os from kyber import Kyber512, Kyber768, Kyber1024 -from aes256_ctr_drbg import AES256_CTR_DRBG +from drbg.aes256_ctr_drbg import AES256_CTR_DRBG def parse_kat_data(data): diff --git a/test_kyber_kat.py b/tests/test_kyber_kat.py similarity index 90% rename from test_kyber_kat.py rename to tests/test_kyber_kat.py index 8809899..f5a3f8b 100644 --- a/test_kyber_kat.py +++ b/tests/test_kyber_kat.py @@ -1,6 +1,12 @@ +""" +An alternative way of checking the Kyber KAT +Does nothing which isn't already checked in test_kyber.py +""" + + from kyber import Kyber512, Kyber768, Kyber1024 from hashlib import sha256 -from aes256_ctr_drbg import AES256_CTR_DRBG +from drbg.aes256_ctr_drbg import AES256_CTR_DRBG def generate_kat_hash(kyber): diff --git a/tests/test_ml_kem.py b/tests/test_ml_kem.py new file mode 100644 index 0000000..24aa5fd --- /dev/null +++ b/tests/test_ml_kem.py @@ -0,0 +1,30 @@ +import unittest +from ml_kem import ML_KEM128, ML_KEM192, ML_KEM256 + +class TestML_KEM(unittest.TestCase): + """ + Test ML_KEM levels for internal + consistency by generating key pairs + and shared secrets. + """ + + def generic_test_ML_KEM(self, ML_KEM, count): + for _ in range(count): + (ek, dk) = ML_KEM.keygen() + for _ in range(count): + (K, c) = ML_KEM.encaps(ek) + K_prime = ML_KEM.decaps(c, dk) + self.assertEqual(K, K_prime) + + def test_ML_KEM128(self): + self.generic_test_ML_KEM(ML_KEM128, 5) + + def test_ML_KEM192(self): + self.generic_test_ML_KEM(ML_KEM192, 5) + + def test_ML_KEM256(self): + self.generic_test_ML_KEM(ML_KEM256, 5) + + +if __name__ == "__main__": + unittest.main() diff --git a/utilities/__init__.py b/utilities/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/utils.py b/utilities/utils.py similarity index 100% rename from utils.py rename to utilities/utils.py