Fix details by hand

docs/asymmetric.rst

@@ -112,4 +112,38 @@ ML-KEM
 >>> 
 >>> ss_recv = mlkem_priv.decapsulate(ct)
 >>> ss_send == ss_recv
+True
+
+ML-DSA
+------
+
+.. autoclass:: MlDsaType
+    :show-inheritance:
+
+.. autoclass:: MlDsaPublic
+    :private-members:
+    :members:
+    :inherited-members:
+    
+.. autoclass:: MlDsaPrivate
+    :members:
+    :inherited-members:
+
+**Example:**
+
+>>> from wolfcrypt.ciphers import MlDsaType, MlDsaPrivate, MlDsaPublic
+>>> 
+>>> mldsa_type = MlDsaType.ML_DSA_44
+>>> 
+>>> mldsa_priv = MlDsaPrivate.make_key(mldsa_type)
+>>> pub_key = mldsa_priv.encode_pub_key()
+>>> 
+>>> mldsa_pub = MlDsaPublic(mldsa_type)
+>>> mldsa_pub.decode_key(pub_key)
+>>> 
+>>> msg = "This is an example message"
+>>> 
+>>> sig = mldsa_priv.sign(msg)
+>>> 
+>>> mldsa_pub.verify(sig, msg)
 True

scripts/build_ffi.py

@@ -374,6 +374,7 @@ def get_features(local_wolfssl, features):
     features["AESGCM_STREAM"] = 1 if '#define WOLFSSL_AESGCM_STREAM' in defines else 0
     features["RSA_PSS"] = 1 if '#define WC_RSA_PSS' in defines else 0
     features["CHACHA20_POLY1305"] = 1 if '#define HAVE_CHACHA' and '#define HAVE_POLY1305' in defines else 0
+    features["ML_DSA"] = 1 if '#define WOLFSSL_WC_DILITHIUM' in defines else 0
 
     if '#define HAVE_FIPS' in defines:
         if not fips:
@@ -935,12 +936,16 @@ def build_ffi(local_wolfssl, features):
         int wolfCrypt_GetPrivateKeyReadEnable_fips(enum wc_KeyType);
         """
 
+    if features["ML_KEM"] or features["ML_DSA"]:
+        cdef += """
+        static const int INVALID_DEVID;
+        """
+
     if features["ML_KEM"]:
         cdef += """
         static const int WC_ML_KEM_512;
         static const int WC_ML_KEM_768;
         static const int WC_ML_KEM_1024;
-        static const int INVALID_DEVID;
         typedef struct {...; } KyberKey;
         int wc_KyberKey_CipherTextSize(KyberKey* key, word32* len);
         int wc_KyberKey_SharedSecretSize(KyberKey* key, word32* len);
@@ -968,9 +973,6 @@ def build_ffi(local_wolfssl, features):
         int wc_dilithium_init_ex(dilithium_key* key, void* heap, int devId);
         int wc_dilithium_set_level(dilithium_key* key, byte level);
         void wc_dilithium_free(dilithium_key* key);
-        int wc_dilithium_priv_size(dilithium_key* key);
-        int wc_dilithium_pub_size(dilithium_key* key);
-        int wc_dilithium_sig_size(dilithium_key* key);
         int wc_dilithium_make_key(dilithium_key* key, WC_RNG* rng);
         int wc_dilithium_export_private(dilithium_key* key, byte* out, word32* outLen);
         int wc_dilithium_import_private(const byte* priv, word32 privSz, dilithium_key* key);
@@ -978,6 +980,10 @@ def build_ffi(local_wolfssl, features):
         int wc_dilithium_import_public(const byte* in, word32 inLen, dilithium_key* key);
         int wc_dilithium_sign_msg(const byte* msg, word32 msgLen, byte* sig, word32* sigLen, dilithium_key* key, WC_RNG* rng);
         int wc_dilithium_verify_msg(const byte* sig, word32 sigLen, const byte* msg, word32 msgLen, int* res, dilithium_key* key);
+        typedef dilithium_key MlDsaKey;
+        int wc_MlDsaKey_GetPrivLen(MlDsaKey* key, int* len);
+        int wc_MlDsaKey_GetPubLen(MlDsaKey* key, int* len);
+        int wc_MlDsaKey_GetSigLen(MlDsaKey* key, int* len);
         """
 
     ffibuilder.cdef(cdef)

tests/test_mldsa.py

@@ -22,19 +22,19 @@
 
 from wolfcrypt._ffi import lib as _lib
 
-if hasattr(_lib, "ML_DSA_ENABLED") and _lib.ML_DSA_ENABLED:
+if _lib.ML_DSA_ENABLED:
-    from binascii import unhexlify as h2b
-
     import pytest
 
-    from wolfcrypt.mldsa import MlDsaPrivate, MlDsaPublic, MlDsaType
+    from wolfcrypt.ciphers import MlDsaPrivate, MlDsaPublic, MlDsaType
     from wolfcrypt.random import Random
 
     @pytest.fixture
     def rng():
         return Random()
 
-    @pytest.fixture(params=[MlDsaType.ML_DSA_44, MlDsaType.ML_DSA_65, MlDsaType.ML_DSA_87])
+    @pytest.fixture(
+        params=[MlDsaType.ML_DSA_44, MlDsaType.ML_DSA_65, MlDsaType.ML_DSA_87]
+    )
     def mldsa_type(request):
         return request.param
 
@@ -45,71 +45,75 @@ if hasattr(_lib, "ML_DSA_ENABLED") and _lib.ML_DSA_ENABLED:
         mldsa_pub = MlDsaPublic(mldsa_type)
         assert isinstance(mldsa_pub, MlDsaPublic)
 
-    def test_key_sizes(mldsa_type):
+    def test_size_properties(mldsa_type):
-        mldsa_priv = MlDsaPrivate(mldsa_type)
+        refvals = {
+            MlDsaType.ML_DSA_44: {
+                "sig_size": 2420,
+                "pub_key_size": 1312,
+                "priv_key_size": 2560,
+            },
+            MlDsaType.ML_DSA_65: {
+                "sig_size": 3309,
+                "pub_key_size": 1952,
+                "priv_key_size": 4032,
+            },
+            MlDsaType.ML_DSA_87: {
+                "sig_size": 4627,
+                "pub_key_size": 2592,
+                "priv_key_size": 4896,
+            },
+        }
 
-        # Check that key sizes are returned correctly
-        assert mldsa_priv.priv_key_size > 0
-        assert mldsa_priv.pub_key_size > 0
-        assert mldsa_priv.sig_size > 0
-
-        # Public key should have the same pub_key_size
         mldsa_pub = MlDsaPublic(mldsa_type)
-        assert mldsa_pub.pub_key_size == mldsa_priv.pub_key_size
+        assert mldsa_pub.sig_size == refvals[mldsa_type]["sig_size"]
-        assert mldsa_pub.sig_size == mldsa_priv.sig_size
+        assert mldsa_pub.key_size == refvals[mldsa_type]["pub_key_size"]
 
-    """
+        mldsa_priv = MlDsaPrivate(mldsa_type)
-    def test_key_generation(mldsa_type, rng):
+        assert mldsa_priv.sig_size == refvals[mldsa_type]["sig_size"]
-        # Test key generation
+        assert mldsa_priv.pub_key_size == refvals[mldsa_type]["pub_key_size"]
+        assert mldsa_priv.priv_key_size == refvals[mldsa_type]["priv_key_size"]
+
+    def test_initializations(mldsa_type, rng):
         mldsa_priv = MlDsaPrivate.make_key(mldsa_type, rng)
-        assert isinstance(mldsa_priv, MlDsaPrivate)
+        assert type(mldsa_priv) is MlDsaPrivate
 
-        # Export keys
+        mldsa_priv2 = MlDsaPrivate(mldsa_type)
+        assert type(mldsa_priv2) is MlDsaPrivate
+
+        mldsa_pub = MlDsaPublic(mldsa_type)
+        assert type(mldsa_pub) is MlDsaPublic
+
+    def test_key_import_export(mldsa_type, rng):
+        # Generate key pair and export keys
+        mldsa_priv = MlDsaPrivate.make_key(mldsa_type, rng)
         priv_key = mldsa_priv.encode_priv_key()
         pub_key = mldsa_priv.encode_pub_key()
-
-        # Check key sizes
         assert len(priv_key) == mldsa_priv.priv_key_size
         assert len(pub_key) == mldsa_priv.pub_key_size
-    """
 
-    """
+        # Export key pair from imported one
-    def test_key_import_export(mldsa_type, rng):
-        # Generate a key pair
-        mldsa_priv = MlDsaPrivate.make_key(mldsa_type, rng)
-
-        # Export keys
-        priv_key = mldsa_priv.encode_priv_key()
-        pub_key = mldsa_priv.encode_pub_key()
-
-        # Import private key
         mldsa_priv2 = MlDsaPrivate(mldsa_type)
-        mldsa_priv2.decode_key(priv_key)
+        mldsa_priv2.decode_key(priv_key, pub_key)
-
-        # Export keys from imported private key
         priv_key2 = mldsa_priv2.encode_priv_key()
         pub_key2 = mldsa_priv2.encode_pub_key()
-
-        # Keys should match
         assert priv_key == priv_key2
         assert pub_key == pub_key2
 
-        # Import public key
+        # Export private key from imported one
+        mldsa_priv3 = MlDsaPrivate(mldsa_type)
+        mldsa_priv3.decode_key(priv_key)
+        priv_key3 = mldsa_priv3.encode_priv_key()
+        assert priv_key == priv_key3
+
+        # Export public key from imported one
         mldsa_pub = MlDsaPublic(mldsa_type)
         mldsa_pub.decode_key(pub_key)
-
-        # Export public key from imported public key
         pub_key3 = mldsa_pub.encode_key()
-
-        # Public keys should match
         assert pub_key == pub_key3
-    """
 
     def test_sign_verify(mldsa_type, rng):
-        # Generate a key pair
+        # Generate a key pair and export public key
         mldsa_priv = MlDsaPrivate.make_key(mldsa_type, rng)
-
-        # Export public key
         pub_key = mldsa_priv.encode_pub_key()
 
         # Import public key
@@ -119,32 +123,14 @@ if hasattr(_lib, "ML_DSA_ENABLED") and _lib.ML_DSA_ENABLED:
         # Sign a message
         message = b"This is a test message for ML-DSA signature"
         signature = mldsa_priv.sign(message, rng)
+        assert len(signature) == mldsa_priv.sig_size
 
-        # Verify the signature
+        # Verify the signature by MlDsaPrivate
+        assert mldsa_priv.verify(signature, message)
+
+        # Verify the signature by MlDsaPublic
         assert mldsa_pub.verify(signature, message)
 
         # Verify with wrong message
         wrong_message = b"This is a wrong message for ML-DSA signature"
         assert not mldsa_pub.verify(signature, wrong_message)
-
-    """
-    def test_der_encoding(mldsa_type, rng):
-        # Generate a key pair
-        mldsa_priv = MlDsaPrivate.make_key(mldsa_type, rng)
-
-        # Export keys in DER format
-        priv_key_der = mldsa_priv.encode_priv_key_der()
-        pub_key_der = mldsa_priv.encode_pub_key_der()
-
-        # Check that DER encoded keys are longer than raw keys
-        assert len(priv_key_der) > mldsa_priv.priv_key_size
-        assert len(pub_key_der) > mldsa_priv.pub_key_size
-
-        # Test public key DER encoding from public key object
-        mldsa_pub = MlDsaPublic(mldsa_type)
-        mldsa_pub.decode_key(mldsa_priv.encode_pub_key())
-        pub_key_der2 = mldsa_pub.encode_key_der()
-
-        # DER encoded public keys should match
-        assert pub_key_der == pub_key_der2
-    """

wolfcrypt/ciphers.py

@@ -20,6 +20,8 @@
 
 # pylint: disable=no-member,no-name-in-module
 
+from enum import IntEnum
+
 from wolfcrypt._ffi import ffi as _ffi
 from wolfcrypt._ffi import lib as _lib
 from wolfcrypt.utils import t2b
@@ -1671,8 +1673,6 @@ if _lib.ED448_ENABLED:
 
 
 if _lib.ML_KEM_ENABLED:
-    from enum import IntEnum
-
     class MlKemType(IntEnum):
         """
         `MlKemType` specifies supported ML-KEM types.
@@ -1953,3 +1953,277 @@ if _lib.ML_KEM_ENABLED:
                 raise WolfCryptError("wc_KyberKey_Decapsulate() error (%d)" % ret)
 
             return _ffi.buffer(ss, ss_size)[:]
+
+
+if _lib.ML_DSA_ENABLED:
+    class MlDsaType(IntEnum):
+        """
+        `MlDsaType` specifies supported ML-DSA types.
+
+        `MlDsaType` is arguments for constructors and some initialization functions for `MlDsaPublic` and `MlDsaPrivate`.
+
+        Followings are all possible values:
+
+        - `ML_DSA_44`
+        - `ML_DSA_65`
+        - `ML_DSA_87`
+        """
+
+        ML_DSA_44 = _lib.WC_ML_DSA_44
+        ML_DSA_65 = _lib.WC_ML_DSA_65
+        ML_DSA_87 = _lib.WC_ML_DSA_87
+
+    class _MlDsaBase(object):
+        INVALID_DEVID = _lib.INVALID_DEVID
+
+        def __init__(self, mldsa_type):
+            self._init_done = False
+            self.native_object = _ffi.new("dilithium_key *")
+            ret = _lib.wc_dilithium_init_ex(
+                self.native_object, _ffi.NULL, self.INVALID_DEVID
+            )
+
+            if ret < 0:  # pragma: no cover
+                raise WolfCryptError("wc_dilithium_init_ex() error (%d)" % ret)
+
+            self._init_done = True
+
+            ret = _lib.wc_dilithium_set_level(self.native_object, mldsa_type)
+
+            if ret < 0:  # pragma: no cover
+                raise WolfCryptError("wc_dilithium_set_level() error (%d)" % ret)
+
+        def __del__(self):
+            if self._init_done:
+                _lib.wc_dilithium_free(self.native_object)
+
+        @property
+        def _pub_key_size(self):
+            size = _ffi.new("int *")
+            ret = _lib.wc_MlDsaKey_GetPubLen(self.native_object, size)
+
+            if ret < 0:  # pragma: no cover
+                raise WolfCryptError("wc_MlDsaKey_GetPubLen() error (%d)" % ret)
+
+            return size[0]
+
+        @property
+        def sig_size(self):
+            """
+            :return: signature size in bytes
+            :rtype: int
+            """
+            size = _ffi.new("int *")
+            ret = _lib.wc_MlDsaKey_GetSigLen(self.native_object, size)
+
+            if ret < 0:  # pragma: no cover
+                raise WolfCryptError("wc_MlDsaKey_GetSigLen() error (%d)" % ret)
+
+            return size[0]
+
+        def _decode_pub_key(self, pub_key):
+            pub_key_bytestype = t2b(pub_key)
+            ret = _lib.wc_dilithium_import_public(
+                _ffi.from_buffer(pub_key_bytestype),
+                len(pub_key_bytestype),
+                self.native_object,
+            )
+
+            if ret < 0:  # pragma: no cover
+                raise WolfCryptError("wc_dilithium_import_public() error (%d)" % ret)
+
+        def _encode_pub_key(self):
+            in_size = self._pub_key_size
+            pub_key = _ffi.new(f"byte[{in_size}]")
+            out_size = _ffi.new("word32 *")
+            out_size[0] = in_size
+            ret = _lib.wc_dilithium_export_public(self.native_object, pub_key, out_size)
+
+            if ret < 0:  # pragma: no cover
+                raise WolfCryptError("wc_dilithium_export_public() error (%d)" % ret)
+
+            if in_size != out_size[0]:
+                raise WolfCryptError(
+                    "in_size=%d and out_size=%d don't match" % (in_size, out_size[0])
+                )
+
+            return _ffi.buffer(pub_key, out_size[0])[:]
+
+        def verify(self, signature, message):
+            """
+            :param signature: signature to be verified
+            :type signature: bytes or str
+            :param message: message to be verified
+            :type message: bytes or str
+            :return: True if the verification is successful, False otherwise
+            :rtype: bool
+            """
+            sig_bytestype = t2b(signature)
+            msg_bytestype = t2b(message)
+            res = _ffi.new("int *")
+
+            ret = _lib.wc_dilithium_verify_msg(
+                _ffi.from_buffer(sig_bytestype),
+                len(sig_bytestype),
+                _ffi.from_buffer(msg_bytestype),
+                len(msg_bytestype),
+                res,
+                self.native_object,
+            )
+
+            if ret < 0:  # pragma: no cover
+                raise WolfCryptError("wc_dilithium_verify_msg() error (%d)" % ret)
+
+            return res[0] == 1
+
+    class MlDsaPrivate(_MlDsaBase):
+        @classmethod
+        def make_key(cls, mldsa_type, rng=Random()):
+            """
+            :param mldsa_type: ML-DSA type
+            :type mldsa_type: MlDsaType
+            :param rng: random number generator for a key generation
+            :type rng: Random
+            :return: `MlDsaPrivate` object
+            :rtype: MlDsaPrivate
+            """
+            mldsa_priv = cls(mldsa_type)
+            ret = _lib.wc_dilithium_make_key(
+                mldsa_priv.native_object, rng.native_object
+            )
+
+            if ret < 0:  # pragma: no cover
+                raise WolfCryptError("wc_dilithium_make_key() error (%d)" % ret)
+
+            return mldsa_priv
+
+        @property
+        def pub_key_size(self):
+            """
+            :return: public key size in bytes
+            :rtype: int
+            """
+            return self._pub_key_size
+
+        @property
+        def priv_key_size(self):
+            """
+            :return: private key size in bytes
+            :rtype: int
+            """
+            size = _ffi.new("int *")
+            ret = _lib.wc_MlDsaKey_GetPrivLen(self.native_object, size)
+
+            if ret < 0:  # pragma: no cover
+                raise WolfCryptError("wc_MlDsaKey_GetPrivLen() error (%d)" % ret)
+
+            key_pair_size = size[0]
+
+            return key_pair_size - self.pub_key_size
+
+        def encode_pub_key(self):
+            """
+            :return: exported public key
+            :rtype: bytes
+            """
+            return self._encode_pub_key()
+
+        def encode_priv_key(self):
+            """
+            :return: exported private key
+            :rtype: bytes
+            """
+            in_size = self.priv_key_size
+            priv_key = _ffi.new(f"byte[{in_size}]")
+            out_size = _ffi.new("word32 *")
+            out_size[0] = in_size
+            ret = _lib.wc_dilithium_export_private(
+                self.native_object, priv_key, out_size
+            )
+
+            if ret < 0:  # pragma: no cover
+                raise WolfCryptError("wc_dilithium_export_private() error (%d)" % ret)
+
+            if in_size != out_size[0]:
+                raise WolfCryptError(
+                    "in_size=%d and out_size=%d don't match" % (in_size, out_size[0])
+                )
+
+            return _ffi.buffer(priv_key, out_size[0])[:]
+
+        def decode_key(self, priv_key, pub_key=None):
+            """
+            :param priv_key: private key to be imported
+            :type priv_key: bytes or str
+            :param pub_key: public key to be imported
+            :type pub_key: bytes or str or None
+            """
+            priv_key_bytestype = t2b(priv_key)
+            ret = _lib.wc_dilithium_import_private(
+                _ffi.from_buffer(priv_key_bytestype),
+                len(priv_key_bytestype),
+                self.native_object,
+            )
+
+            if ret < 0:  # pragma: no cover
+                raise WolfCryptError("wc_dilithium_import_private() error (%d)" % ret)
+
+            if pub_key is not None:
+                self._decode_pub_key(pub_key)
+
+        def sign(self, message, rng=Random()):
+            """
+            :param message: message to be signed
+            :type message: bytes or str
+            :param rng: random number generator for sign
+            :type rng: Random
+            :return: signature
+            :rtype: bytes
+            """
+            msg_bytestype = t2b(message)
+            in_size = self.sig_size
+            signature = _ffi.new(f"byte[{in_size}]")
+            out_size = _ffi.new("word32 *")
+            out_size[0] = in_size
+
+            ret = _lib.wc_dilithium_sign_msg(
+                _ffi.from_buffer(msg_bytestype),
+                len(msg_bytestype),
+                signature,
+                out_size,
+                self.native_object,
+                rng.native_object,
+            )
+
+            if ret < 0:  # pragma: no cover
+                raise WolfCryptError("wc_dilithium_sign_msg() error (%d)" % ret)
+
+            if in_size != out_size[0]:
+                raise WolfCryptError(
+                    "in_size=%d and out_size=%d don't match" % (in_size, out_size[0])
+                )
+
+            return _ffi.buffer(signature, out_size[0])[:]
+
+    class MlDsaPublic(_MlDsaBase):
+        @property
+        def key_size(self):
+            """
+            :return: public key size in bytes
+            :rtype: int
+            """
+            return self._pub_key_size
+
+        def decode_key(self, pub_key):
+            """
+            :param pub_key: public key to be imported
+            :type pub_key: bytes or str
+            """
+            return self._decode_pub_key(pub_key)
+
+        def encode_key(self):
+            """
+            :return: exported public key
+            :rtype: bytes
+            """
+            return self._encode_pub_key()

wolfcrypt/mldsa.py

@@ -1,312 +0,0 @@
-# mldsa.py
-#
-# Copyright (C) 2025 wolfSSL Inc.
-#
-# This file is part of wolfSSL. (formerly known as CyaSSL)
-#
-# wolfSSL is free software; you can redistribute it and/or modify
-# it under the terms of the GNU General Public License as published by
-# the Free Software Foundation; either version 2 of the License, or
-# (at your option) any later version.
-#
-# wolfSSL is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-# GNU General Public License for more details.
-#
-# You should have received a copy of the GNU General Public License
-# along with this program; if not, write to the Free Software
-# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA
-
-from enum import IntEnum
-
-from wolfcrypt._ffi import ffi as _ffi
-from wolfcrypt._ffi import lib as _lib
-from wolfcrypt.utils import t2b
-from wolfcrypt.random import Random
-from wolfcrypt.exceptions import WolfCryptError
-
-if hasattr(_lib, "wc_dilithium_init_ex"):
-    class MlDsaType(IntEnum):
-        """
-        `MlDsaType` specifies supported ML-DSA types.
-
-        `MlDsaType` is arguments for constructors and some initialization functions for `MlDsaPublic` and `MlDsaPrivate`.
-
-        Followings are all possible values:
-
-        - `ML_DSA_44`
-        - `ML_DSA_65`
-        - `ML_DSA_87`
-        """
-
-        ML_DSA_44 = _lib.WC_ML_DSA_44
-        ML_DSA_65 = _lib.WC_ML_DSA_65
-        ML_DSA_87 = _lib.WC_ML_DSA_87
-
-    class _MlDsaBase(object):
-        INVALID_DEVID = _lib.INVALID_DEVID
-
-        def __init__(self, mldsa_type):
-            self.init_done = False
-            self.native_object = _ffi.new("dilithium_key *")
-            ret = _lib.wc_dilithium_init_ex(
-                self.native_object, _ffi.NULL, self.INVALID_DEVID
-            )
-
-            if ret < 0:  # pragma: no cover
-                raise WolfCryptError("wc_dilithium_init_ex() error (%d)" % ret)
-
-            ret = _lib.wc_dilithium_set_level(self.native_object, mldsa_type)
-
-            if ret < 0:  # pragma: no cover
-                raise WolfCryptError("wc_dilithium_set_level() error (%d)" % ret)
-
-            self.init_done = True
-
-        def __del__(self):
-            if self.init_done:
-                _lib.wc_dilithium_free(self.native_object)
-
-        @property
-        def priv_key_size(self):
-            """
-            :return: private key size in bytes
-            :rtype: int
-            """
-            ret = _lib.wc_dilithium_priv_size(self.native_object)
-
-            if ret < 0:  # pragma: no cover
-                raise WolfCryptError("wc_dilithium_priv_size() error (%d)" % ret)
-
-            return ret
-
-        @property
-        def pub_key_size(self):
-            """
-            :return: public key size in bytes
-            :rtype: int
-            """
-            ret = _lib.wc_dilithium_pub_size(self.native_object)
-
-            if ret < 0:  # pragma: no cover
-                raise WolfCryptError("wc_dilithium_pub_size() error (%d)" % ret)
-
-            return ret
-
-        @property
-        def sig_size(self):
-            """
-            :return: signature size in bytes
-            :rtype: int
-            """
-            ret = _lib.wc_dilithium_sig_size(self.native_object)
-
-            if ret < 0:  # pragma: no cover
-                raise WolfCryptError("wc_dilithium_sig_size() error (%d)" % ret)
-
-            return ret
-
-        def _encode_pub_key(self):
-            pub_key_size = self.pub_key_size
-            pub_key = _ffi.new(f"unsigned char[{pub_key_size}]")
-            out_len = _ffi.new("word32 *")
-            out_len[0] = pub_key_size
-            ret = _lib.wc_dilithium_export_public(
-                self.native_object, pub_key, out_len
-            )
-
-            if ret < 0:  # pragma: no cover
-                raise WolfCryptError("wc_dilithium_export_public() error (%d)" % ret)
-
-            return _ffi.buffer(pub_key, out_len[0])[:]
-
-        def _encode_pub_key_der(self, with_alg=1):
-            pub_key_size = self.pub_key_size
-            # DER encoding adds some overhead
-            der_size = pub_key_size + 50
-            output = _ffi.new(f"unsigned char[{der_size}]")
-            # Use export_public instead of PublicKeyToDer since there's no direct equivalent
-            out_len = _ffi.new("word32 *")
-            out_len[0] = der_size
-            ret = _lib.wc_dilithium_export_public(
-                self.native_object, output, out_len
-            )
-
-            if ret < 0:  # pragma: no cover
-                raise WolfCryptError("wc_dilithium_export_public() error (%d)" % ret)
-
-            return _ffi.buffer(output, ret)[:]
-
-    class MlDsaPrivate(_MlDsaBase):
-        @classmethod
-        def make_key(cls, mldsa_type, rng=Random()):
-            """
-            :param mldsa_type: ML-DSA type
-            :type mldsa_type: MlDsaType
-            :param rng: random number generator for a key generation
-            :type rng: Random
-            :return: `MlDsaPrivate` object
-            :rtype: MlDsaPrivate
-            """
-            mldsa_priv = cls(mldsa_type)
-            ret = _lib.wc_dilithium_make_key(mldsa_priv.native_object, rng.native_object)
-
-            if ret < 0:  # pragma: no cover
-                raise WolfCryptError("wc_dilithium_make_key() error (%d)" % ret)
-
-            return mldsa_priv
-
-        def decode_key(self, priv_key):
-            """
-            :param priv_key: private key to be imported
-            :type priv_key: bytes or str
-            """
-            priv_key_bytestype = t2b(priv_key)
-            ret = _lib.wc_dilithium_import_private(
-                _ffi.from_buffer(priv_key_bytestype),
-                len(priv_key_bytestype),
-                self.native_object
-            )
-
-            if ret < 0:  # pragma: no cover
-                raise WolfCryptError("wc_dilithium_import_private() error (%d)" % ret)
-
-        def encode_pub_key(self):
-            """
-            :return: exported public key
-            :rtype: bytes
-            """
-            return self._encode_pub_key()
-
-        def encode_pub_key_der(self, with_alg=1):
-            """
-            :return: exported public key in DER format
-            :rtype: bytes
-            """
-            return self._encode_pub_key_der(with_alg)
-
-        def encode_priv_key(self):
-            """
-            :return: exported private key
-            :rtype: bytes
-            """
-            priv_key_size = self.priv_key_size
-            priv_key = _ffi.new(f"unsigned char[{priv_key_size}]")
-            out_len = _ffi.new("word32 *")
-            out_len[0] = priv_key_size
-            ret = _lib.wc_dilithium_export_private(
-                self.native_object, priv_key, out_len
-            )
-
-            if ret < 0:  # pragma: no cover
-                raise WolfCryptError("wc_dilithium_export_private() error (%d)" % ret)
-
-            return _ffi.buffer(priv_key, out_len[0])[:]
-
-        def encode_priv_key_der(self):
-            """
-            :return: exported private key in DER format
-            :rtype: bytes
-            """
-            priv_key_size = self.priv_key_size
-            # DER encoding adds some overhead
-            der_size = priv_key_size + 50
-            output = _ffi.new(f"unsigned char[{der_size}]")
-            # Use export_private instead of PrivateKeyToDer since there's no direct equivalent
-            out_len = _ffi.new("word32 *")
-            out_len[0] = der_size
-            ret = _lib.wc_dilithium_export_private(
-                self.native_object, output, out_len
-            )
-
-            if ret < 0:  # pragma: no cover
-                raise WolfCryptError("wc_dilithium_export_private() error (%d)" % ret)
-
-            return _ffi.buffer(output, out_len[0])[:]
-
-        def sign(self, message, rng=Random()):
-            """
-            :param message: message to be signed
-            :type message: bytes or str
-            :param rng: random number generator
-            :type rng: Random
-            :return: signature
-            :rtype: bytes
-            """
-            message = t2b(message)
-            sig_size = self.sig_size
-            signature = _ffi.new(f"unsigned char[{sig_size}]")
-            sig_len = _ffi.new("word32 *")
-            sig_len[0] = sig_size
-
-            ret = _lib.wc_dilithium_sign_msg(
-                _ffi.from_buffer(message),
-                len(message),
-                signature,
-                sig_len,
-                self.native_object,
-                rng.native_object
-            )
-
-            if ret < 0:  # pragma: no cover
-                raise WolfCryptError("wc_dilithium_sign_msg() error (%d)" % ret)
-
-            return _ffi.buffer(signature, sig_len[0])[:] 
-
-    class MlDsaPublic(_MlDsaBase):
-        def decode_key(self, pub_key):
-            """
-            :param pub_key: public key to be imported
-            :type pub_key: bytes or str
-            """
-            pub_key_bytestype = t2b(pub_key)
-            ret = _lib.wc_dilithium_import_public(
-                _ffi.from_buffer(pub_key_bytestype),
-                len(pub_key_bytestype),
-                self.native_object
-            )
-
-            if ret < 0:  # pragma: no cover
-                raise WolfCryptError("wc_dilithium_import_public() error (%d)" % ret)
-
-        def encode_key(self):
-            """
-            :return: exported public key
-            :rtype: bytes
-            """
-            return self._encode_pub_key()
-
-        def encode_key_der(self, with_alg=1):
-            """
-            :return: exported public key in DER format
-            :rtype: bytes
-            """
-            return self._encode_pub_key_der(with_alg)
-
-        def verify(self, signature, message):
-            """
-            :param signature: signature to be verified
-            :type signature: bytes or str
-            :param message: message to be verified
-            :type message: bytes or str
-            :return: True if the signature is valid, False otherwise
-            :rtype: bool
-            """
-            signature = t2b(signature)
-            message = t2b(message)
-            res = _ffi.new("int *")
-
-            ret = _lib.wc_dilithium_verify_msg(
-                _ffi.from_buffer(signature),
-                len(signature),
-                _ffi.from_buffer(message),
-                len(message),
-                res,
-                self.native_object
-            )
-
-            if ret < 0:  # pragma: no cover
-                raise WolfCryptError("wc_dilithium_verify_msg() error (%d)" % ret)
-
-            return res[0] == 1