Skip to content

Add support for ML-DSA "post-quantum" key types #721

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
dralley wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Add support for ML-DSA "post-quantum" key types #721

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
81 changes: 74 additions & 7 deletions src/trustme/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,9 +11,16 @@
from typing import TYPE_CHECKING, Generator, List, Optional, Union

import idna
import cryptography
from cryptography import x509
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import ec, rsa

_MLDSA_AVAILABLE = tuple(
int(x) for x in cryptography.__version__.split(".")[:2]
) >= (49, 0)
if _MLDSA_AVAILABLE:
from cryptography.hazmat.primitives.asymmetric import mldsa
from cryptography.hazmat.primitives.serialization import (
Encoding,
NoEncryption,
Expand All @@ -26,9 +33,22 @@

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]
from cryptography.hazmat.primitives.asymmetric import mldsa as _mldsa_types

CERTIFICATE_PUBLIC_KEY_TYPES = Union[
rsa.RSAPublicKey,
ec.EllipticCurvePublicKey,
_mldsa_types.MLDSA44PublicKey,
_mldsa_types.MLDSA65PublicKey,
_mldsa_types.MLDSA87PublicKey,
]
CERTIFICATE_PRIVATE_KEY_TYPES = Union[
rsa.RSAPrivateKey,
ec.EllipticCurvePrivateKey,
_mldsa_types.MLDSA44PrivateKey,
_mldsa_types.MLDSA65PrivateKey,
_mldsa_types.MLDSA87PrivateKey,
]

__all__ = ["CA"]

Expand Down Expand Up @@ -214,6 +234,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:
Expand All @@ -223,9 +246,49 @@ 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 in (KeyType.MLDSA44, KeyType.MLDSA65, KeyType.MLDSA87):
if not _MLDSA_AVAILABLE:
raise TypeError(
f"{self.name} requires cryptography >= 49 with ML-DSA support"
)
cls = {
KeyType.MLDSA44: mldsa.MLDSA44PrivateKey,
KeyType.MLDSA65: mldsa.MLDSA65PrivateKey,
KeyType.MLDSA87: mldsa.MLDSA87PrivateKey,
}[self]
return cls.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 _MLDSA_AVAILABLE:
if 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
raise TypeError(f"Unsupported key type: {type(private_key)}")


class CA:
"""A certificate authority."""
Expand All @@ -241,6 +304,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

Expand All @@ -250,9 +314,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(
Expand Down Expand Up @@ -286,7 +352,7 @@ def __init__(
critical=True,
).sign(
private_key=sign_key,
algorithm=hashes.SHA256(),
algorithm=sign_key_type._hash_algorithm,
)

@property
Expand All @@ -301,7 +367,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()
)
)

Expand Down Expand Up @@ -440,7 +506,7 @@ def issue_cert(
)
.sign(
private_key=self._private_key,
algorithm=hashes.SHA256(),
algorithm=self._key_type._hash_algorithm,
)
)

Expand All @@ -453,7 +519,7 @@ def issue_cert(
return LeafCert(
key.private_bytes(
Encoding.PEM,
PrivateFormat.TraditionalOpenSSL,
key_type._private_key_format,
NoEncryption(),
),
cert.public_bytes(Encoding.PEM),
Expand Down Expand Up @@ -499,6 +565,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

Expand Down
31 changes: 26 additions & 5 deletions tests/test_trustme.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,13 @@
SslSocket = Union[ssl.SSLSocket, OpenSSL.SSL.Connection]


from trustme import _MLDSA_AVAILABLE

_skip_mldsa = pytest.mark.skipif(
not _MLDSA_AVAILABLE, reason="cryptography lacks ML-DSA support"
)


def _path_length(ca_cert: x509.Certificate) -> Optional[int]:
bc = ca_cert.extensions.get_extension_for_class(x509.BasicConstraints)
return bc.value.path_length
Expand Down Expand Up @@ -66,16 +73,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"),
pytest.param(KeyType.MLDSA44, b"BEGIN PRIVATE KEY", marks=_skip_mldsa),
pytest.param(KeyType.MLDSA65, b"BEGIN PRIVATE KEY", marks=_skip_mldsa),
pytest.param(KeyType.MLDSA87, b"BEGIN PRIVATE KEY", marks=_skip_mldsa),
],
)
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)
Expand Down Expand Up @@ -357,7 +369,16 @@ 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,
pytest.param(KeyType.MLDSA44, marks=_skip_mldsa),
pytest.param(KeyType.MLDSA65, marks=_skip_mldsa),
pytest.param(KeyType.MLDSA87, marks=_skip_mldsa),
],
)
def test_stdlib_end_to_end(key_type: KeyType) -> None:
def wrap_client(
ca: CA, raw_client_sock: socket.socket, hostname: str
Expand Down
Loading