Add ML-DSA (FIPS 204) post-quantum key type support

Add MLDSA44, MLDSA65, and MLDSA87 key types for generating
post-quantum certificates using the cryptography library's ML-DSA
support (requires OpenSSL 3.5+).

Author: dralley

src/trustme/__init__.py

@@ -13,7 +13,7 @@
 import idna
 from cryptography import x509
 from cryptography.hazmat.primitives import hashes
-from cryptography.hazmat.primitives.asymmetric import ec, rsa
+from cryptography.hazmat.primitives.asymmetric import ec, mldsa, rsa
 from cryptography.hazmat.primitives.serialization import (
     Encoding,
     NoEncryption,
@@ -27,8 +27,20 @@
 if TYPE_CHECKING:  # pragma: no cover
     import OpenSSL.SSL
 
-    CERTIFICATE_PUBLIC_KEY_TYPES = Union[rsa.RSAPublicKey, ec.EllipticCurvePublicKey]
-    CERTIFICATE_PRIVATE_KEY_TYPES = Union[rsa.RSAPrivateKey, ec.EllipticCurvePrivateKey]
+    CERTIFICATE_PUBLIC_KEY_TYPES = Union[
+        rsa.RSAPublicKey,
+        ec.EllipticCurvePublicKey,
+        mldsa.MLDSA44PublicKey,
+        mldsa.MLDSA65PublicKey,
+        mldsa.MLDSA87PublicKey,
+    ]
+    CERTIFICATE_PRIVATE_KEY_TYPES = Union[
+        rsa.RSAPrivateKey,
+        ec.EllipticCurvePrivateKey,
+        mldsa.MLDSA44PrivateKey,
+        mldsa.MLDSA65PrivateKey,
+        mldsa.MLDSA87PrivateKey,
+    ]
 
 __all__ = ["CA"]
 
@@ -214,6 +226,9 @@ class KeyType(Enum):
 
     RSA = 0
     ECDSA = 1
+    MLDSA44 = 2
+    MLDSA65 = 3
+    MLDSA87 = 4
 
     def _generate_key(self) -> CERTIFICATE_PRIVATE_KEY_TYPES:
         if self is KeyType.RSA:
@@ -223,9 +238,44 @@ def _generate_key(self) -> CERTIFICATE_PRIVATE_KEY_TYPES:
             return rsa.generate_private_key(public_exponent=65537, key_size=2048)
         elif self is KeyType.ECDSA:
             return ec.generate_private_key(ec.SECP256R1())
+        elif self is KeyType.MLDSA44:
+            return mldsa.MLDSA44PrivateKey.generate()
+        elif self is KeyType.MLDSA65:
+            return mldsa.MLDSA65PrivateKey.generate()
+        elif self is KeyType.MLDSA87:
+            return mldsa.MLDSA87PrivateKey.generate()
         else:  # pragma: no cover
             raise ValueError("Unknown key type")
 
+    @property
+    def _hash_algorithm(self) -> Optional[hashes.SHA256]:
+        """ML-DSA uses intrinsic hashing; RSA/ECDSA use SHA-256."""
+        if self in (KeyType.MLDSA44, KeyType.MLDSA65, KeyType.MLDSA87):
+            return None
+        return hashes.SHA256()
+
+    @property
+    def _private_key_format(self) -> PrivateFormat:
+        """ML-DSA keys don't support TraditionalOpenSSL format."""
+        if self in (KeyType.MLDSA44, KeyType.MLDSA65, KeyType.MLDSA87):
+            return PrivateFormat.PKCS8
+        return PrivateFormat.TraditionalOpenSSL
+
+
+def _detect_key_type(private_key: CERTIFICATE_PRIVATE_KEY_TYPES) -> KeyType:
+    if isinstance(private_key, rsa.RSAPrivateKey):
+        return KeyType.RSA
+    elif isinstance(private_key, ec.EllipticCurvePrivateKey):
+        return KeyType.ECDSA
+    elif isinstance(private_key, mldsa.MLDSA44PrivateKey):
+        return KeyType.MLDSA44
+    elif isinstance(private_key, mldsa.MLDSA65PrivateKey):
+        return KeyType.MLDSA65
+    elif isinstance(private_key, mldsa.MLDSA87PrivateKey):
+        return KeyType.MLDSA87
+    else:
+        raise TypeError(f"Unsupported key type: {type(private_key)}")
+
 
 class CA:
     """A certificate authority."""
@@ -241,6 +291,7 @@ def __init__(
         key_type: KeyType = KeyType.ECDSA,
     ) -> None:
         self.parent_cert = parent_cert
+        self._key_type = key_type
         self._private_key = key_type._generate_key()
         self._path_length = path_length
 
@@ -250,9 +301,11 @@ def __init__(
         )
         issuer = name
         sign_key = self._private_key
+        sign_key_type = key_type
         aki: Optional[x509.AuthorityKeyIdentifier]
         if parent_cert is not None:
             sign_key = parent_cert._private_key
+            sign_key_type = parent_cert._key_type
             parent_certificate = parent_cert._certificate
             issuer = parent_certificate.subject
             ski_ext = parent_certificate.extensions.get_extension_for_class(
@@ -286,7 +339,7 @@ def __init__(
             critical=True,
         ).sign(
             private_key=sign_key,
-            algorithm=hashes.SHA256(),
+            algorithm=sign_key_type._hash_algorithm,
         )
 
     @property
@@ -301,7 +354,7 @@ def private_key_pem(self) -> Blob:
         other certificates from this CA."""
         return Blob(
             self._private_key.private_bytes(
-                Encoding.PEM, PrivateFormat.TraditionalOpenSSL, NoEncryption()
+                Encoding.PEM, self._key_type._private_key_format, NoEncryption()
             )
         )
 
@@ -440,7 +493,7 @@ def issue_cert(
             )
             .sign(
                 private_key=self._private_key,
-                algorithm=hashes.SHA256(),
+                algorithm=self._key_type._hash_algorithm,
             )
         )
 
@@ -453,7 +506,7 @@ def issue_cert(
         return LeafCert(
             key.private_bytes(
                 Encoding.PEM,
-                PrivateFormat.TraditionalOpenSSL,
+                key_type._private_key_format,
                 NoEncryption(),
             ),
             cert.public_bytes(Encoding.PEM),
@@ -499,6 +552,7 @@ def from_pem(cls, cert_bytes: bytes, private_key_bytes: bytes) -> "CA":
         ca.parent_cert = None
         ca._certificate = x509.load_pem_x509_certificate(cert_bytes)
         ca._private_key = load_pem_private_key(private_key_bytes, password=None)  # type: ignore[assignment]
+        ca._key_type = _detect_key_type(ca._private_key)
 
         return ca
 

tests/test_trustme.py

@@ -66,16 +66,21 @@ def assert_is_leaf(leaf_cert: x509.Certificate) -> None:
 
 
 @pytest.mark.parametrize(
-    "key_type,expected_key_header", [(KeyType.RSA, b"RSA"), (KeyType.ECDSA, b"EC")]
+    "key_type,expected_key_header",
+    [
+        (KeyType.RSA, b"BEGIN RSA PRIVATE KEY"),
+        (KeyType.ECDSA, b"BEGIN EC PRIVATE KEY"),
+        (KeyType.MLDSA44, b"BEGIN PRIVATE KEY"),
+        (KeyType.MLDSA65, b"BEGIN PRIVATE KEY"),
+        (KeyType.MLDSA87, b"BEGIN PRIVATE KEY"),
+    ],
 )
 def test_basics(key_type: KeyType, expected_key_header: bytes) -> None:
     ca = CA(key_type=key_type)
 
     today = datetime.datetime.now(datetime.timezone.utc)
 
-    assert (
-        b"BEGIN " + expected_key_header + b" PRIVATE KEY" in ca.private_key_pem.bytes()
-    )
+    assert expected_key_header in ca.private_key_pem.bytes()
     assert b"BEGIN CERTIFICATE" in ca.cert_pem.bytes()
 
     private_key = load_pem_private_key(ca.private_key_pem.bytes(), password=None)
@@ -357,7 +362,10 @@ def doit(ca: CA, hostname: str, server_cert: LeafCert) -> None:
         doit(bad_ca, hostname, ca.issue_cert(hostname, key_type=key_type))
 
 
-@pytest.mark.parametrize("key_type", [KeyType.RSA, KeyType.ECDSA])
+@pytest.mark.parametrize(
+    "key_type",
+    [KeyType.RSA, KeyType.ECDSA, KeyType.MLDSA44, KeyType.MLDSA65, KeyType.MLDSA87],
+)
 def test_stdlib_end_to_end(key_type: KeyType) -> None:
     def wrap_client(
         ca: CA, raw_client_sock: socket.socket, hostname: str