mlkem 0.0.1__tar.gz

mlkem-0.0.1/PKG-INFO

@@ -0,0 +1,125 @@
+Metadata-Version: 2.2
+Name: mlkem
+Version: 0.0.1
+Summary: A module-latice-based key-encapsulation mechanism (ML-KEM). Implements the FIPS-203 standard.
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+
+# Module-Lattice-Based Key-Encapsulation Mechanism (ML-KEM)
+An implementation of the module-lattice-based key encapsulation mechanism (ML-KEM)
+as described in [FIPS-203](https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.203.pdf).
+At this time the package is in alpha and _SHOULD NOT_ be considered for real-world
+cryptographic applications.
+
+# Usage
+
+The package includes includes a pure python implementation of the K-PKE function
+(`mlkem.k_pke.K_PKE`) and an implementation that leverages C extensions
+(`mlkem.fast_k_pke.Fast_K_PKE`). The implementations have interchangeable interfaces
+and can be selected in their wrapper class `mlkem.ml_kem.ML_KEM` by setting the
+`fast` param to `True` for C extensions and `False` for pure python e.g.
+
+```python
+from mlkem.ml_kem import ML_KEM
+from mlkem.parameter_set import ParameterSet
+
+ML_KEM(ParameterSet.ML_KEM_768, fast=True)  # C extensions
+ML_KEM(ParameterSet.ML_KEM_768, fast=False)  # Pure python
+```
+
+Both implementations are self contained and portable (assuming you have 8 bits per byte
+on your system) with no dependencies on third party libraries in either the C or python
+code.
+
+NIST recommends the ML-KEM-768 parameter set, which offers 192 bit security. ML-KEM-512
+and ML-KEM-1024 are also available, which provide 128 and 256 bit security respectively.
+ML-KEM-768 is used by default in this package. Thus, the two instantiations below are
+equivalent -
+
+```python
+from mlkem.ml_kem import ML_KEM
+from mlkem.parameter_set import ParameterSet
+
+ML_KEM()
+ML_KEM(ParameterSet.ML_KEM_768, fast=True)
+```
+
+The interface follows the one defined in section 7 of the standard for the functions KeyGen,
+Encaps and Decaps.
+
+```python
+from mlkem.ml_kem import ML_KEM
+
+ml_kem = ML_KEM()
+ek, dk = ml_kem.key_gen()  # encapsulation and decapsulation key
+k, c = ml_kem.encaps(ek)  # shared secret key and ciphertext
+k_ = ml_kem.decaps(dk, c)  # shared secret key
+```
+
+In a less contrived scenario, Alice might run KeyGen and send the encapsulation key
+to Bob. Bob would then run Encaps and generate a shared secret key and a ciphertext.
+Bob would send the ciphertext to Alice, who would derive the shared secret key from the
+ciphertext. Alice and Bob can then use the shared secret key to generate additional
+secret material by passing it to a KDF, use the shared secret to directly key a symmetric
+cipher like AES, etc.
+
+# Development
+
+As a prerequisite, `uv` is required for this project
+
+    pip install uv
+
+Build the C extensions
+
+    uv run python setup.py build_ext --inplace
+
+Run the test suite
+
+    uv run pytest
+
+Build the docs
+
+    uv run make -C docs html
+
+# Performance
+
+The performance of the C extensions is _significantly_ faster. The python
+implementation is primarily included for those that wish to explore and debug the algorithm.
+Performance against the NIST
+[keygen](https://github.com/usnistgov/ACVP-Server/tree/master/gen-val/json-files/ML-KEM-keyGen-FIPS203) and
+[encap/decap](https://github.com/usnistgov/ACVP-Server/tree/master/gen-val/json-files/ML-KEM-encapDecap-FIPS203)
+test vectors can be seen below -
+
+#### C Extensions
+```
+uv run pytest -k "key_gen or encaps or decaps"
+================================================ test session starts ================================================
+platform darwin -- Python 3.11.11, pytest-8.3.4, pluggy-1.5.0
+rootdir: /Users/antonku/dev/github/mlkem
+configfile: pyproject.toml
+plugins: cov-6.0.0
+collected 261 items / 36 deselected / 225 selected
+
+tests/test_decaps.py ...........................................................................              [ 33%]
+tests/test_encaps.py ...........................................................................              [ 66%]
+tests/test_key_gen.py ...........................................................................             [100%]
+
+======================================== 225 passed, 36 deselected in 0.23s =========================================
+```
+
+#### Pure Python
+```
+uv run pytest -k "key_gen or encaps or decaps"                                                              1 ↵
+================================================ test session starts ================================================
+platform darwin -- Python 3.11.11, pytest-8.3.4, pluggy-1.5.0
+rootdir: /Users/antonku/dev/github/mlkem
+configfile: pyproject.toml
+plugins: cov-6.0.0
+collected 261 items / 36 deselected / 225 selected
+
+tests/test_decaps.py ...........................................................................              [ 33%]
+tests/test_encaps.py ...........................................................................              [ 66%]
+tests/test_key_gen.py ...........................................................................             [100%]
+
+======================================== 225 passed, 36 deselected in 4.42s =========================================
+```

mlkem-0.0.1/README.md

@@ -0,0 +1,118 @@
+# Module-Lattice-Based Key-Encapsulation Mechanism (ML-KEM)
+An implementation of the module-lattice-based key encapsulation mechanism (ML-KEM)
+as described in [FIPS-203](https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.203.pdf).
+At this time the package is in alpha and _SHOULD NOT_ be considered for real-world
+cryptographic applications.
+
+# Usage
+
+The package includes includes a pure python implementation of the K-PKE function
+(`mlkem.k_pke.K_PKE`) and an implementation that leverages C extensions
+(`mlkem.fast_k_pke.Fast_K_PKE`). The implementations have interchangeable interfaces
+and can be selected in their wrapper class `mlkem.ml_kem.ML_KEM` by setting the
+`fast` param to `True` for C extensions and `False` for pure python e.g.
+
+```python
+from mlkem.ml_kem import ML_KEM
+from mlkem.parameter_set import ParameterSet
+
+ML_KEM(ParameterSet.ML_KEM_768, fast=True)  # C extensions
+ML_KEM(ParameterSet.ML_KEM_768, fast=False)  # Pure python
+```
+
+Both implementations are self contained and portable (assuming you have 8 bits per byte
+on your system) with no dependencies on third party libraries in either the C or python
+code.
+
+NIST recommends the ML-KEM-768 parameter set, which offers 192 bit security. ML-KEM-512
+and ML-KEM-1024 are also available, which provide 128 and 256 bit security respectively.
+ML-KEM-768 is used by default in this package. Thus, the two instantiations below are
+equivalent -
+
+```python
+from mlkem.ml_kem import ML_KEM
+from mlkem.parameter_set import ParameterSet
+
+ML_KEM()
+ML_KEM(ParameterSet.ML_KEM_768, fast=True)
+```
+
+The interface follows the one defined in section 7 of the standard for the functions KeyGen,
+Encaps and Decaps.
+
+```python
+from mlkem.ml_kem import ML_KEM
+
+ml_kem = ML_KEM()
+ek, dk = ml_kem.key_gen()  # encapsulation and decapsulation key
+k, c = ml_kem.encaps(ek)  # shared secret key and ciphertext
+k_ = ml_kem.decaps(dk, c)  # shared secret key
+```
+
+In a less contrived scenario, Alice might run KeyGen and send the encapsulation key
+to Bob. Bob would then run Encaps and generate a shared secret key and a ciphertext.
+Bob would send the ciphertext to Alice, who would derive the shared secret key from the
+ciphertext. Alice and Bob can then use the shared secret key to generate additional
+secret material by passing it to a KDF, use the shared secret to directly key a symmetric
+cipher like AES, etc.
+
+# Development
+
+As a prerequisite, `uv` is required for this project
+
+    pip install uv
+
+Build the C extensions
+
+    uv run python setup.py build_ext --inplace
+
+Run the test suite
+
+    uv run pytest
+
+Build the docs
+
+    uv run make -C docs html
+
+# Performance
+
+The performance of the C extensions is _significantly_ faster. The python
+implementation is primarily included for those that wish to explore and debug the algorithm.
+Performance against the NIST
+[keygen](https://github.com/usnistgov/ACVP-Server/tree/master/gen-val/json-files/ML-KEM-keyGen-FIPS203) and
+[encap/decap](https://github.com/usnistgov/ACVP-Server/tree/master/gen-val/json-files/ML-KEM-encapDecap-FIPS203)
+test vectors can be seen below -
+
+#### C Extensions
+```
+uv run pytest -k "key_gen or encaps or decaps"
+================================================ test session starts ================================================
+platform darwin -- Python 3.11.11, pytest-8.3.4, pluggy-1.5.0
+rootdir: /Users/antonku/dev/github/mlkem
+configfile: pyproject.toml
+plugins: cov-6.0.0
+collected 261 items / 36 deselected / 225 selected
+
+tests/test_decaps.py ...........................................................................              [ 33%]
+tests/test_encaps.py ...........................................................................              [ 66%]
+tests/test_key_gen.py ...........................................................................             [100%]
+
+======================================== 225 passed, 36 deselected in 0.23s =========================================
+```
+
+#### Pure Python
+```
+uv run pytest -k "key_gen or encaps or decaps"                                                              1 ↵
+================================================ test session starts ================================================
+platform darwin -- Python 3.11.11, pytest-8.3.4, pluggy-1.5.0
+rootdir: /Users/antonku/dev/github/mlkem
+configfile: pyproject.toml
+plugins: cov-6.0.0
+collected 261 items / 36 deselected / 225 selected
+
+tests/test_decaps.py ...........................................................................              [ 33%]
+tests/test_encaps.py ...........................................................................              [ 66%]
+tests/test_key_gen.py ...........................................................................             [100%]
+
+======================================== 225 passed, 36 deselected in 4.42s =========================================
+```

mlkem-0.0.1/mlkem/auxiliary/crypto.py

@@ -0,0 +1,48 @@
+from hashlib import sha3_256, sha3_512, shake_128, shake_256
+
+
+def prf(eta: int, s: bytes, b: bytes) -> bytes:
+    if eta not in {2, 3}:
+        raise ValueError(f"eta must be 2 or 3 (got {eta})")
+    if len(s) != 32:
+        raise ValueError(f"len(s) must be 32 (got {len(s)})")
+    if len(b) != 1:
+        raise ValueError(f"len(b) must be 1 (got {len(b)})")
+
+    # length passed to digest is byte length, so omit factor of 8 from spec (which uses bit length)
+    return shake_256(s + b).digest(64 * eta)
+
+
+def h(s: bytes) -> bytes:
+    return sha3_256(s).digest()
+
+
+def j(s: bytes) -> bytes:
+    # length passed to digest is byte length, so omit factor of 8 from spec (which uses bit length)
+    return shake_256(s).digest(32)
+
+
+def g(c: bytes) -> tuple[bytes, bytes]:
+    ab = sha3_512(c).digest()
+    return ab[:32], ab[32:]
+
+
+class XOF:
+    def __init__(self) -> None:
+        # https://cryptojedi.org/papers/terminate-20230516.pdf
+        self.chunk_size = 840
+        self.shake = shake_128()
+        self.data = b""
+        self.idx = 0
+
+    def absorb(self, string: bytes) -> None:
+        self.shake.update(string)
+        self.data += self.shake.digest(self.chunk_size)
+
+    def squeeze(self, length: int) -> bytes:
+        while self.idx + length > len(self.data):
+            self.data += self.shake.digest(self.chunk_size)
+
+        result = self.data[self.idx : self.idx + length]
+        self.idx += length
+        return result

mlkem-0.0.1/mlkem/auxiliary/general.py

@@ -0,0 +1,112 @@
+from mlkem.math.constants import n, q
+from mlkem.math.field import Zm
+
+BITS_IN_BYTE = 8
+MAX_D = q.bit_length()
+
+
+def bits_to_bytes(bits: list[int]) -> list[int]:
+    """Converts a bit array (of a length that is a multiple of 8) into an array of bytes.
+
+    Bytes are represented as unsigned numbers in the range [0, 255]. Bits are either 0 or 1.
+
+    Args:
+        | bits (:type:`list[int]`): The bit array (of a length that is a multiple of 8).
+
+    Returns:
+        :type:`list[int]`: The array of bytes equivalent to the bit array.
+    """
+    length = len(bits)
+    if length % BITS_IN_BYTE != 0:
+        raise ValueError(
+            f"Bit array must have a length that is a multiple of 8 (got {length})."
+        )
+
+    result = [0 for _ in range(length // BITS_IN_BYTE)]
+    for i in range(length):
+        bitval = bits[i] * (1 << (i % BITS_IN_BYTE))
+        result[i // BITS_IN_BYTE] = result[i // BITS_IN_BYTE] + bitval
+
+    return result
+
+
+def bytes_to_bits(byts: list[int]) -> list[int]:
+    """Converts a byte array into an array of bits.
+
+    Bytes are represented as unsigned numbers in the range [0, 255]. Bits are either 0 or 1.
+
+    Args:
+        | byts (:type:`list[int]`): The byte array.
+
+    Returns:
+        :type:`list[int]`: The array of bits equivalent to the byte array.
+    """
+    c = byts.copy()
+
+    result: list[int] = []
+    for i in range(len(c)):
+        for _ in range(BITS_IN_BYTE):
+            result.append(c[i] & 1)
+            c[i] //= 2
+
+    return result
+
+
+def _round_fraction(x: int, y: int) -> int:
+    """Round the fraction x/y to the nearest integer."""
+    return (2 * x + y) // (2 * y)
+
+
+def compress(d: int, x: Zm) -> Zm:
+    if not d < MAX_D:
+        raise ValueError(f"d must be less than {MAX_D} (got {d}).")
+    if x.m != q:
+        raise ValueError(f"Element being compressed must be in Z_q (got Z_{x.m}).")
+
+    m = 1 << d
+    val = _round_fraction(m * x.val, q) % m
+    return Zm(val, m)
+
+
+def decompress(d: int, y: Zm) -> Zm:
+    if not d < MAX_D:
+        raise ValueError(f"d must be less than {MAX_D} (got {d}).")
+
+    m = 1 << d
+    val = _round_fraction(q * y.val, m)
+    return Zm(val, q)
+
+
+def byte_encode(d: int, f: list[Zm]) -> bytes:
+    if len(f) != n:
+        raise ValueError(f"f must have {n} elements (got {len(f)}).")
+
+    if d > MAX_D or d < 1:
+        raise ValueError(f"d may not be greater than {MAX_D} or less than 1 (got {d}).")
+
+    b = [0 for _ in range(n * d)]
+    for i in range(n):
+        a = f[i].val
+
+        for j in range(d):
+            x = a & 1
+            b[i * d + j] = x
+            a = (a - x) // 2
+
+    return bytes(bits_to_bytes(b))
+
+
+def byte_decode(d: int, b: bytes) -> list[Zm]:
+    if d > MAX_D or d < 1:
+        raise ValueError(f"d may not be greater than {MAX_D} or less than 1 (got {d}).")
+
+    m = q if d == MAX_D else 1 << d
+    bits = bytes_to_bits([x for x in b])
+
+    f = []
+    for i in range(n):
+        fi = sum([bits[i * d + j] * (1 << j) for j in range(d)])
+        fi_m = Zm(fi, m)
+        f.append(fi_m)
+
+    return f

mlkem-0.0.1/mlkem/auxiliary/ntt.py

@@ -0,0 +1,128 @@
+from mlkem.math.constants import n, q
+from mlkem.math.field import Zm
+from mlkem.math.polynomial_ring import PolynomialRing, RingRepresentation
+
+# see Appendix A - Precomputed Values for the NTT
+ZETA_LOOKUP = [
+    Zm(x, q) for x in
+    [
+        1, 1729, 2580, 3289, 2642, 630, 1897, 848,
+        1062, 1919, 193, 797, 2786, 3260, 569, 1746,
+        296, 2447, 1339, 1476, 3046, 56, 2240, 1333,
+        1426, 2094, 535, 2882, 2393, 2879, 1974, 821,
+        289, 331, 3253, 1756, 1197, 2304, 2277, 2055,
+        650, 1977, 2513, 632, 2865, 33, 1320, 1915,
+        2319, 1435, 807, 452, 1438, 2868, 1534, 2402,
+        2647, 2617, 1481, 648, 2474, 3110, 1227, 910,
+        17, 2761, 583, 2649, 1637, 723, 2288, 1100,
+        1409, 2662, 3281, 233, 756, 2156, 3015, 3050,
+        1703, 1651, 2789, 1789, 1847, 952, 1461, 2687,
+        939, 2308, 2437, 2388, 733, 2337, 268, 641,
+        1584, 2298, 2037, 3220, 375, 2549, 2090, 1645,
+        1063, 319, 2773, 757, 2099, 561, 2466, 2594,
+        2804, 1092, 403, 1026, 1143, 2150, 2775, 886,
+        1722, 1212, 1874, 1029, 2110, 2935, 885, 2154
+    ]
+]  # fmt: skip
+
+GAMMA_LOOKUP = [
+    Zm(x, q) for x in
+    [
+        17, -17, 2761, -2761, 583, -583, 2649, -2649,
+        1637, -1637, 723, -723, 2288, -2288, 1100, -1100,
+        1409, -1409, 2662, -2662, 3281, -3281, 233, -233,
+        756, -756, 2156, -2156, 3015, -3015, 3050, -3050,
+        1703, -1703, 1651, -1651, 2789, -2789, 1789, -1789,
+        1847, -1847, 952, -952, 1461, -1461, 2687, -2687,
+        939, -939, 2308, -2308, 2437, -2437, 2388, -2388,
+        733, -733, 2337, -2337, 268, -268, 641, -641,
+        1584, -1584, 2298, -2298, 2037, -2037, 3220, -3220,
+        375, -375, 2549, -2549, 2090, -2090, 1645, -1645,
+        1063, -1063, 319, -319, 2773, -2773, 757, -757,
+        2099, -2099, 561, -561, 2466, -2466, 2594, -2594,
+        2804, -2804, 1092, -1092, 403, -403, 1026, -1026,
+        1143, -1143, 2150, -2150, 2775, -2775, 886, -886,
+        1722, -1722, 1212, -1212, 1874, -1874, 1029, -1029,
+        2110, -2110, 2935, -2935, 885, -885, 2154, -2154
+    ]
+]  # fmt: skip
+
+
+def ntt(f: PolynomialRing) -> PolynomialRing:
+    if f.representation != RingRepresentation.STANDARD:
+        raise ValueError(
+            "NTT can only be applied to polynomials in standard representation."
+        )
+
+    f_ = PolynomialRing(f.coefficients, RingRepresentation.NTT)
+    i = 1
+    length = 128
+
+    while length >= 2:
+        for start in range(0, n, 2 * length):
+            zeta = ZETA_LOOKUP[i]
+            i += 1
+
+            for j in range(start, start + length):
+                t = zeta * f_[j + length]
+                f_[j + length] = f_[j] - t
+                f_[j] += t
+
+        length //= 2
+
+    return f_
+
+
+def ntt_inv(f_: PolynomialRing) -> PolynomialRing:
+    if f_.representation != RingRepresentation.NTT:
+        raise ValueError(
+            "Inverse NTT can only be applied to polynomials in NTT representation."
+        )
+
+    f = PolynomialRing(f_.coefficients, RingRepresentation.STANDARD)
+    i = 127
+    length = 2
+
+    while length <= 128:
+        for start in range(0, n, 2 * length):
+            zeta = ZETA_LOOKUP[i]
+            i -= 1
+
+            for j in range(start, start + length):
+                t = f[j]
+                f[j] = t + f[j + length]
+                f[j + length] = zeta * (f[j + length] - t)
+
+        length *= 2
+
+    scale = Zm(3303, q)  # 3303 = 128^{-1} mod q
+    for i, fi in enumerate(f.coefficients):
+        f[i] = fi * scale
+
+    return f
+
+
+def multiply_ntt(f_: PolynomialRing, g_: PolynomialRing) -> PolynomialRing:
+    if (
+        f_.representation != RingRepresentation.NTT
+        or g_.representation != RingRepresentation.NTT
+    ):
+        raise ValueError(
+            "NTT multiplication can only be applied to polynomials in NTT representation."
+        )
+
+    h_ = PolynomialRing(representation=RingRepresentation.NTT)
+
+    for i in range(128):
+        gamma = GAMMA_LOOKUP[i]
+        h_[2 * i], h_[2 * i + 1] = _base_case_multiply(
+            f_[2 * i], f_[2 * i + 1], g_[2 * i], g_[2 * i + 1], gamma
+        )
+
+    return h_
+
+
+def _base_case_multiply(a0: Zm, a1: Zm, b0: Zm, b1: Zm, gamma: Zm) -> tuple[Zm, Zm]:
+    c0 = a0 * b0 + a1 * b1 * gamma
+    c1 = a0 * b1 + a1 * b0
+    return c0, c1

mlkem-0.0.1/mlkem/auxiliary/sampling.py

@@ -0,0 +1,47 @@
+from mlkem.auxiliary.crypto import XOF
+from mlkem.auxiliary.general import bytes_to_bits
+from mlkem.math.constants import n, q
+from mlkem.math.field import Zm
+from mlkem.math.polynomial_ring import PolynomialRing, RingRepresentation
+
+
+def sample_ntt(b: bytes) -> PolynomialRing:
+    if len(b) != 34:
+        raise ValueError(
+            f"Input must be 34 bytes (32-byte seed and two indices). Got {len(b)}."
+        )
+
+    a = PolynomialRing(representation=RingRepresentation.NTT)
+    xof = XOF()
+    xof.absorb(b)
+
+    j = 0
+    while j < n:
+        c = xof.squeeze(3)
+        d1 = c[0] + n * (c[1] % 16)
+        d2 = c[1] // 16 + 16 * c[2]
+
+        if d1 < q:
+            a[j] = Zm(d1, q)
+            j += 1
+
+        if d2 < q and j < n:
+            a[j] = Zm(d2, q)
+            j += 1
+
+    return a
+
+
+def sample_poly_cbd(eta: int, b: bytes) -> PolynomialRing:
+    if len(b) != 64 * eta:
+        raise ValueError(f"Input must be {64 * eta} bytes (got {len(b)}).")
+
+    f = PolynomialRing(representation=RingRepresentation.STANDARD)
+    bits = bytes_to_bits([x for x in b])
+
+    for i in range(n):
+        x = sum([bits[2 * i * eta + j] for j in range(eta)])
+        y = sum([bits[2 * i * eta + eta + j] for j in range(eta)])
+        f[i] = Zm(x - y, q)
+
+    return f

mlkem-0.0.1/mlkem/data_types.py

@@ -0,0 +1,18 @@
+from abc import abstractmethod
+from typing import Protocol, Self
+
+
+# the entries of a matrix must have a type supporting addition
+# define that interface via this protocol
+class Field(Protocol):
+    @abstractmethod
+    def __add__(self, other: Self) -> Self:
+        pass
+
+    @abstractmethod
+    def __sub__(self, other: Self) -> Self:
+        pass
+
+    @abstractmethod
+    def __mul__(self, other: Self) -> Self:
+        pass