diff --git a/src/trustme/__init__.py b/src/trustme/__init__.py index 1040892..8bbc30b 100644 --- a/src/trustme/__init__.py +++ b/src/trustme/__init__.py @@ -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, @@ -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"] @@ -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: @@ -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.""" @@ -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 @@ -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( @@ -286,7 +352,7 @@ def __init__( critical=True, ).sign( private_key=sign_key, - algorithm=hashes.SHA256(), + algorithm=sign_key_type._hash_algorithm, ) @property @@ -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() ) ) @@ -440,7 +506,7 @@ def issue_cert( ) .sign( private_key=self._private_key, - algorithm=hashes.SHA256(), + algorithm=self._key_type._hash_algorithm, ) ) @@ -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), @@ -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 diff --git a/tests/test_trustme.py b/tests/test_trustme.py index 93275bc..5330c4c 100644 --- a/tests/test_trustme.py +++ b/tests/test_trustme.py @@ -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 @@ -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) @@ -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