dilithium-py 1.0.0__py3-none-any.whl

dilithium_py/dilithium/__init__.py

@@ -0,0 +1,3 @@
+from .default_parameters import Dilithium2, Dilithium3, Dilithium5
+
+__all__ = ["Dilithium2", "Dilithium3", "Dilithium5"]

dilithium_py/dilithium/default_parameters.py

@@ -0,0 +1,38 @@
+from .dilithium import Dilithium
+
+DEFAULT_PARAMETERS = {
+    "dilithium2": {
+        "d": 13,
+        "k": 4,
+        "l": 4,
+        "eta": 2,
+        "tau": 39,
+        "omega": 80,
+        "gamma_1": 131072,  # 2^17
+        "gamma_2": 95232,  # (q-1)/88
+    },
+    "dilithium3": {
+        "d": 13,
+        "k": 6,
+        "l": 5,
+        "eta": 4,
+        "tau": 49,
+        "omega": 55,
+        "gamma_1": 524288,  # 2^19
+        "gamma_2": 261888,  # (q-1)/32
+    },
+    "dilithium5": {
+        "d": 13,
+        "k": 8,
+        "l": 7,
+        "eta": 2,
+        "tau": 60,
+        "omega": 75,
+        "gamma_1": 524288,  # 2^19
+        "gamma_2": 261888,  # (q-1)/32
+    },
+}
+
+Dilithium2 = Dilithium(DEFAULT_PARAMETERS["dilithium2"])
+Dilithium3 = Dilithium(DEFAULT_PARAMETERS["dilithium3"])
+Dilithium5 = Dilithium(DEFAULT_PARAMETERS["dilithium5"])

dilithium_py/dilithium/dilithium.py

@@ -0,0 +1,312 @@
+import os
+from ..modules.modules import ModuleDilithium
+
+try:
+    from xoflib import shake256
+except ImportError:
+    from ..shake.shake_wrapper import shake256
+
+
+class Dilithium:
+    def __init__(self, parameter_set):
+        self.d = parameter_set["d"]
+        self.k = parameter_set["k"]
+        self.l = parameter_set["l"]
+        self.eta = parameter_set["eta"]
+        self.tau = parameter_set["tau"]
+        self.omega = parameter_set["omega"]
+        self.gamma_1 = parameter_set["gamma_1"]
+        self.gamma_2 = parameter_set["gamma_2"]
+        self.beta = self.tau * self.eta
+
+        self.M = ModuleDilithium()
+        self.R = self.M.ring
+
+        # Use system randomness by default, for deterministic randomness
+        # use the method `set_drbg_seed()`
+        self.random_bytes = os.urandom
+
+    def set_drbg_seed(self, seed):
+        """
+        Change entropy source to a DRBG and seed it with provided value.
+
+        Setting the seed switches the entropy source from :func:`os.urandom()`
+        to an AES256 CTR DRBG.
+
+        Used for both deterministic versions of Kyber as well as testing
+        alignment with the KAT vectors
+
+        Note:
+          currently requires pycryptodome for AES impl.
+        """
+        try:
+            from ..drbg.aes256_ctr_drbg import AES256_CTR_DRBG
+
+            self._drbg = AES256_CTR_DRBG(seed)
+            self.random_bytes = self._drbg.random_bytes
+        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"
+            )
+
+    """
+    H() uses Shake256 to hash data to 32 and 64 bytes in a 
+    few places in the code 
+    """
+
+    @staticmethod
+    def _h(input_bytes, length):
+        """
+        H: B^*  -> B^*
+        """
+        return shake256(input_bytes).read(length)
+
+    def _expand_matrix_from_seed(self, rho):
+        """
+        Helper function which generates a element of size
+        k x l from a seed `rho`.
+        """
+        A_data = [[0 for _ in range(self.l)] for _ in range(self.k)]
+        for i in range(self.k):
+            for j in range(self.l):
+                A_data[i][j] = self.R.rejection_sample_ntt_poly(rho, i, j)
+        return self.M(A_data)
+
+    def _expand_vector_from_seed(self, rho_prime):
+        s1_elements = [
+            self.R.rejection_bounded_poly(rho_prime, i, self.eta) for i in range(self.l)
+        ]
+        s2_elements = [
+            self.R.rejection_bounded_poly(rho_prime, i, self.eta)
+            for i in range(self.l, self.l + self.k)
+        ]
+
+        s1 = self.M.vector(s1_elements)
+        s2 = self.M.vector(s2_elements)
+        return s1, s2
+
+    def _expand_mask_vector(self, rho_prime, kappa):
+        elements = [
+            self.R.sample_mask_polynomial(rho_prime, i, kappa, self.gamma_1)
+            for i in range(self.l)
+        ]
+        return self.M.vector(elements)
+
+    @staticmethod
+    def _pack_pk(rho, t1):
+        return rho + t1.bit_pack_t1()
+
+    def _pack_sk(self, rho, K, tr, s1, s2, t0):
+        s1_bytes = s1.bit_pack_s(self.eta)
+        s2_bytes = s2.bit_pack_s(self.eta)
+        t0_bytes = t0.bit_pack_t0()
+        return rho + K + tr + s1_bytes + s2_bytes + t0_bytes
+
+    def _pack_h(self, h):
+        non_zero_positions = [
+            [i for i, c in enumerate(poly.coeffs) if c == 1]
+            for row in h._data
+            for poly in row
+        ]
+        packed = []
+        offsets = []
+        for positions in non_zero_positions:
+            packed.extend(positions)
+            offsets.append(len(packed))
+
+        padding_len = self.omega - offsets[-1]
+        packed.extend([0 for _ in range(padding_len)])
+        return bytes(packed + offsets)
+
+    def _pack_sig(self, c_tilde, z, h):
+        return c_tilde + z.bit_pack_z(self.gamma_1) + self._pack_h(h)
+
+    def _unpack_pk(self, pk_bytes):
+        rho, t1_bytes = pk_bytes[:32], pk_bytes[32:]
+        t1 = self.M.bit_unpack_t1(t1_bytes, self.k, 1)
+        return rho, t1
+
+    def _unpack_sk(self, sk_bytes):
+        if self.eta == 2:
+            s_bytes = 96
+        else:
+            s_bytes = 128
+        s1_len = s_bytes * self.l
+        s2_len = s_bytes * self.k
+        t0_len = 416 * self.k
+        if len(sk_bytes) != 3 * 32 + s1_len + s2_len + t0_len:
+            raise ValueError("SK packed bytes is of the wrong length")
+
+        # Split bytes between seeds and vectors
+        sk_seed_bytes, sk_vec_bytes = sk_bytes[:96], sk_bytes[96:]
+
+        # Unpack seed bytes
+        rho, K, tr = (
+            sk_seed_bytes[:32],
+            sk_seed_bytes[32:64],
+            sk_seed_bytes[64:96],
+        )
+
+        # Unpack vector bytes
+        s1_bytes = sk_vec_bytes[:s1_len]
+        s2_bytes = sk_vec_bytes[s1_len : s1_len + s2_len]
+        t0_bytes = sk_vec_bytes[-t0_len:]
+
+        # Unpack bytes to vectors
+        s1 = self.M.bit_unpack_s(s1_bytes, self.l, 1, self.eta)
+        s2 = self.M.bit_unpack_s(s2_bytes, self.k, 1, self.eta)
+        t0 = self.M.bit_unpack_t0(t0_bytes, self.k, 1)
+
+        return rho, K, tr, s1, s2, t0
+
+    def _unpack_h(self, h_bytes):
+        offsets = [0] + list(h_bytes[-self.k :])
+        non_zero_positions = [
+            list(h_bytes[offsets[i] : offsets[i + 1]]) for i in range(self.k)
+        ]
+
+        matrix = []
+        for poly_non_zero in non_zero_positions:
+            coeffs = [0 for _ in range(256)]
+            for non_zero in poly_non_zero:
+                coeffs[non_zero] = 1
+            matrix.append([self.R(coeffs)])
+        return self.M(matrix)
+
+    def _unpack_sig(self, sig_bytes):
+        c_tilde = sig_bytes[:32]
+        z_bytes = sig_bytes[32 : -(self.k + self.omega)]
+        h_bytes = sig_bytes[-(self.k + self.omega) :]
+
+        z = self.M.bit_unpack_z(z_bytes, self.l, 1, self.gamma_1)
+        h = self._unpack_h(h_bytes)
+        return c_tilde, z, h
+
+    def keygen(self):
+        """
+        Generates a public-private keyair
+        """
+        # Random seed
+        zeta = self.random_bytes(32)
+
+        # Expand with an XOF (SHAKE256)
+        seed_bytes = self._h(zeta, 128)
+
+        # Split bytes into suitable chunks
+        rho, rho_prime, K = seed_bytes[:32], seed_bytes[32:96], seed_bytes[96:]
+
+        # Generate matrix A ∈ R^(kxl) in the NTT domain
+        A_hat = self._expand_matrix_from_seed(rho)
+
+        # Generate the error vectors s1 ∈ R^l, s2 ∈ R^k
+        s1, s2 = self._expand_vector_from_seed(rho_prime)
+        s1_hat = s1.to_ntt()
+
+        # Matrix multiplication
+        t = (A_hat @ s1_hat).from_ntt() + s2
+
+        t1, t0 = t.power_2_round(self.d)
+
+        # Pack up the bytes
+        pk = self._pack_pk(rho, t1)
+        tr = self._h(pk, 32)
+
+        sk = self._pack_sk(rho, K, tr, s1, s2, t0)
+        return pk, sk
+
+    def sign(self, sk_bytes, m):
+        """
+        Generates a signature for a message m from a byte-encoded private key
+        """
+        # unpack the secret key
+        rho, K, tr, s1, s2, t0 = self._unpack_sk(sk_bytes)
+
+        # Generate matrix A ∈ R^(kxl) in the NTT domain
+        A_hat = self._expand_matrix_from_seed(rho)
+
+        # Set seeds and nonce (kappa)
+        mu = self._h(tr + m, 64)
+        kappa = 0
+        rho_prime = self._h(K + mu, 64)
+
+        # Precompute NTT representation
+        s1 = s1.to_ntt()
+        s2 = s2.to_ntt()
+        t0 = t0.to_ntt()
+
+        alpha = self.gamma_2 << 1
+        while True:
+            y = self._expand_mask_vector(rho_prime, kappa)
+            y_hat = y.to_ntt()
+
+            # increment the nonce
+            kappa += self.l
+
+            w = (A_hat @ y_hat).from_ntt()
+
+            # Extract out both the high and low bits
+            w1, w0 = w.decompose(alpha)
+
+            # Create challenge polynomial
+            w1_bytes = w1.bit_pack_w(self.gamma_2)
+            c_tilde = self._h(mu + w1_bytes, 32)
+            c = self.R.sample_in_ball(c_tilde, self.tau)
+
+            # Store c in NTT form
+            c = c.to_ntt()
+
+            z = y + (s1.scale(c)).from_ntt()
+            if z.check_norm_bound(self.gamma_1 - self.beta):
+                continue
+
+            w0_minus_cs2 = w0 - s2.scale(c).from_ntt()
+            if w0_minus_cs2.check_norm_bound(self.gamma_2 - self.beta):
+                continue
+
+            c_t0 = t0.scale(c).from_ntt()
+            if c_t0.check_norm_bound(self.gamma_2):
+                continue
+
+            w0_minus_cs2_plus_ct0 = w0_minus_cs2 + c_t0
+
+            h = w0_minus_cs2_plus_ct0.make_hint_optimised(w1, alpha)
+            if h.sum_hint() > self.omega:
+                continue
+
+            return self._pack_sig(c_tilde, z, h)
+
+    def verify(self, pk_bytes, m, sig_bytes):
+        """
+        Verifies a signature for a message m from a byte encoded public key and
+        signature
+        """
+        rho, t1 = self._unpack_pk(pk_bytes)
+        c_tilde, z, h = self._unpack_sig(sig_bytes)
+
+        if h.sum_hint() > self.omega:
+            return False
+
+        if z.check_norm_bound(self.gamma_1 - self.beta):
+            return False
+
+        A_hat = self._expand_matrix_from_seed(rho)
+
+        tr = self._h(pk_bytes, 32)
+        mu = self._h(tr + m, 64)
+        c = self.R.sample_in_ball(c_tilde, self.tau)
+
+        # Convert to NTT for computation
+        c = c.to_ntt()
+        z = z.to_ntt()
+
+        t1 = t1.scale(1 << self.d)
+        t1 = t1.to_ntt()
+
+        Az_minus_ct1 = (A_hat @ z) - t1.scale(c)
+        Az_minus_ct1 = Az_minus_ct1.from_ntt()
+
+        w_prime = h.use_hint(Az_minus_ct1, 2 * self.gamma_2)
+        w_prime_bytes = w_prime.bit_pack_w(self.gamma_2)
+
+        return c_tilde == self._h(mu + w_prime_bytes, 32)

dilithium_py/drbg/aes256_ctr_drbg.py

@@ -0,0 +1,134 @@
+import os
+from Crypto.Cipher import AES
+from typing import Optional
+from ..utilities.utils import xor_bytes
+
+
+class AES256_CTR_DRBG:
+    def __init__(self, seed: Optional[bytes] = None, personalization: bytes = b""):
+        """
+        DRBG implementation based on AES-256 CTR following the document NIST SP
+        800-90A Section 10.2.1
+
+        https://csrc.nist.gov/pubs/sp/800/90/a/r1/final
+
+        Used for deterministic randomness, particularly used for comparing the
+        output of Kyber/ML-KEM against known answer tests.
+
+        :param bytes seed: 48 byte seed, if none is supplied a seed is generated
+            using ``os.urandom(48)``.
+        :param bytes personalization: optional bytes, of length at most 48 used
+            during instantiation of the DRBG
+        """
+        self.seed_length = 48
+        self.reseed_interval = 2**48
+        self.key = bytes([0]) * 32
+        self.V = bytes([0]) * 16
+        self.entropy_input = self.__check_entropy_input(seed)
+
+        seed_material = self.__instantiate(personalization=personalization)
+        self.__ctr_drbg_update(seed_material)
+        self.reseed_ctr = 1
+
+    def __check_entropy_input(self, entropy_input: bytes) -> bytes:
+        """
+        If no entropy given, us os.urandom, else
+        check that the input is of the right length.
+        """
+        if entropy_input is None:
+            return os.urandom(self.seed_length)
+        elif len(entropy_input) != self.seed_length:
+            raise ValueError(
+                f"The entropy input must be of length: {self.seed_length}. "
+                f"Input has length {len(entropy_input)}"
+            )
+        return entropy_input
+
+    def __instantiate(self, personalization: bytes = b"") -> bytes:
+        """
+        Combine the input seed and optional personalisation
+        string into the seed material for the DRBG
+
+        Section 10.2.1.3.1, Page 52 (CTR_DRBG_Instantiate_algorithm)
+        """
+        if len(personalization) > self.seed_length:
+            raise ValueError(
+                f"The Personalization String must be at most length: "
+                f"{self.seed_length}. Input has length {len(personalization)}"
+            )
+        # Ensure personalization has exactly seed_length bytes
+        personalization += bytes([0]) * (self.seed_length - len(personalization))
+        # debugging
+        assert len(personalization) == self.seed_length
+        return xor_bytes(self.entropy_input, personalization)
+
+    def __increment_counter(self) -> None:
+        """
+        Increment the internal counter of the DRBG
+        """
+        int_V = int.from_bytes(self.V, "big")
+        new_V = (int_V + 1) % 2**128
+        self.V = new_V.to_bytes(16, byteorder="big")
+
+    def __ctr_drbg_update(self, provided_data: bytes) -> None:
+        """
+        Updates the internal state of the CTR_DRBG using the
+        provided_data
+
+        Section 10.2.1.2, Page 51 (CTR_DRBG_Update)
+        """
+        tmp = b""
+        cipher = AES.new(self.key, AES.MODE_ECB)
+
+        # Collect bytes from AES ECB
+        while len(tmp) != self.seed_length:
+            self.__increment_counter()
+            tmp += cipher.encrypt(self.V)
+
+        # Take the first 48 bytes
+        tmp = tmp[: self.seed_length]
+        tmp = xor_bytes(tmp, provided_data)
+
+        # Set the new values of key and V
+        self.key = tmp[:32]
+        self.V = tmp[32:]
+
+    def random_bytes(self, num_bytes: int, additional: Optional[bytes] = None) -> bytes:
+        """
+        Generate pseudorandom bytes without a generating function
+
+        Section 10.2.1.5.1, Page 56 (CTR_DRBG_Generate_algorithm)
+
+        :param int num_bytes: the number of random bytes requested
+        :param bytes additional: optional bytes to be mixed into the generation
+        :return: pseudorandom bytes extracted from the DRBG of length ``num_bytes``.
+        :rtype: bytes
+        """
+        # 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
+        if additional is None:
+            additional = bytes([0]) * self.seed_length
+        else:
+            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(additional)}"
+                )
+            additional += bytes([0]) * (self.seed_length - len(additional))
+            self.__ctr_drbg_update(additional)
+
+        # Collect bytes!
+        tmp = b""
+        cipher = AES.new(self.key, AES.MODE_ECB)
+        while len(tmp) < num_bytes:
+            self.__increment_counter()
+            tmp += cipher.encrypt(self.V)
+
+        # Collect only the requested number of bits
+        output_bytes = tmp[:num_bytes]
+        self.__ctr_drbg_update(additional)
+        self.reseed_ctr += 1
+        return output_bytes

dilithium_py/ml_dsa/__init__.py

@@ -0,0 +1,3 @@
+from .default_parameters import ML_DSA_44, ML_DSA_65, ML_DSA_87
+
+__all__ = ["ML_DSA_44", "ML_DSA_65", "ML_DSA_87"]

dilithium_py/ml_dsa/default_parameters.py

@@ -0,0 +1,41 @@
+from .ml_dsa import ML_DSA
+
+DEFAULT_PARAMETERS = {
+    "ML_DSA_44": {
+        "d": 13,  # number of bits dropped from t
+        "tau": 39,  # number of ±1 in c
+        "gamma_1": 131072,  # coefficient range of y: 2^17
+        "gamma_2": 95232,  # low order rounding range: (q-1)/88
+        "k": 4,  # Dimensions of A = (k, l)
+        "l": 4,  # Dimensions of A = (k, l)
+        "eta": 2,  # Private key range
+        "omega": 80,  # Max number of ones in hint
+        "c_tilde_bytes": 32,
+    },
+    "ML_DSA_65": {
+        "d": 13,  # number of bits dropped from t
+        "tau": 49,  # number of ±1 in c
+        "gamma_1": 524288,  # coefficient range of y: 2^19
+        "gamma_2": 261888,  # low order rounding range: (q-1)/32
+        "k": 6,  # Dimensions of A = (k, l)
+        "l": 5,  # Dimensions of A = (k, l)
+        "eta": 4,  # Private key range
+        "omega": 55,  # Max number of ones in hint
+        "c_tilde_bytes": 48,
+    },
+    "ML_DSA_87": {
+        "d": 13,  # number of bits dropped from t
+        "tau": 60,  # number of ±1 in c
+        "gamma_1": 524288,  # coefficient range of y: 2^19
+        "gamma_2": 261888,  # low order rounding range: (q-1)/32
+        "k": 8,  # Dimensions of A = (k, l)
+        "l": 7,  # Dimensions of A = (k, l)
+        "eta": 2,  # Private key range
+        "omega": 75,  # Max number of ones in hint
+        "c_tilde_bytes": 64,
+    },
+}
+
+ML_DSA_44 = ML_DSA(DEFAULT_PARAMETERS["ML_DSA_44"])
+ML_DSA_65 = ML_DSA(DEFAULT_PARAMETERS["ML_DSA_65"])
+ML_DSA_87 = ML_DSA(DEFAULT_PARAMETERS["ML_DSA_87"])