pk/tls: Add PrivateKey interface from PEP 543

Author: Synss

src/mbedtls/_tls.pyi

@@ -16,7 +16,6 @@ from mbedtls._tlsi import (
     TLSConfiguration,
     TLSVersion,
 )
-from mbedtls.pk import ECC, RSA
 
 if sys.version_info < (3, 8):
     from typing_extensions import Literal
@@ -84,8 +83,6 @@ class Purpose(enum.IntEnum):
     SERVER_AUTH: int
     CLIENT_AUTH: int
 
-_Key = Union[RSA, ECC]
-PrivateKey = _Key
 CipherSuite = str
 ServerNameCallback = object
 

src/mbedtls/_tlsi.py

@@ -9,15 +9,22 @@
 import os
 import sys
 from dataclasses import dataclass, field
-from typing import Mapping, Optional, Tuple, TypeVar, Union
+from typing import Callable, Mapping, Optional, Tuple, TypeVar, Union
 
 if sys.version_info < (3, 8):
     from typing_extensions import Literal, Protocol
 else:
     from typing import Literal, Protocol
 
+if sys.version_info < (3, 9):
+    _PathLike = os.PathLike
+else:
+    _PathLike = os.PathLike[str]
+
 __all__ = ["NextProtocol", "TLSVersion", "DTLSVersion"]
 
+_Path = Union[_PathLike, str]
+
 
 @enum.unique
 class NextProtocol(enum.Enum):
@@ -59,7 +66,7 @@ def system(cls) -> TrustStore:
         """
 
     @classmethod
-    def from_pem_file(cls, path: Union[str, os.PathLike[str]]) -> TrustStore:
+    def from_pem_file(cls, path: _Path) -> TrustStore:
         """Initializes a trust store from a single file full of PEMs."""
 
 
@@ -78,7 +85,7 @@ def from_buffer(cls, buffer: bytes) -> Certificate:
         """
 
     @classmethod
-    def from_file(cls, path: Union[str, os.PathLike[str]]) -> Certificate:
+    def from_file(cls, path: _Path) -> Certificate:
         """Creates a Certificate object from a file on disk.
 
         This method may be a convenience method that wraps ``open`` and
@@ -89,7 +96,55 @@ def from_file(cls, path: Union[str, os.PathLike[str]]) -> Certificate:
         """
 
 
-PrivateKey = object
+class PrivateKey(Protocol):
+    @classmethod
+    def from_buffer(
+        cls,
+        buffer: bytes,
+        password: Optional[
+            Union[Callable[[], Union[bytes, bytearray]], bytes, bytearray]
+        ] = None,
+    ) -> PrivateKey:
+        """Creates a PrivateKey object from a byte buffer.
+
+        This byte buffer may be either PEM-encoded or DER-encoded. If the
+        buffer is PEM encoded it *must* begin with the standard PEM
+        preamble (a series of dashes followed by the ASCII bytes "BEGIN",
+        the key type, and another series of dashes). In the absence of
+        that preamble, the implementation may assume that the certificate
+        is DER-encoded instead.
+
+        The key may additionally be encrypted. If it is, the ``password``
+        argument can be used to decrypt the key. The ``password`` argument
+        may be a function to call to get the password for decrypting the
+        private key. It will only be called if the private key is encrypted
+        and a password is necessary. It will be called with no arguments,
+        and it should return either bytes or bytearray containing the
+        password. Alternatively a bytes, or bytearray value may be supplied
+        directly as the password argument. It will be ignored if the
+        private key is not encrypted and no password is needed.
+        """
+
+    @classmethod
+    def from_file(
+        cls,
+        path: _Path,
+        password: Optional[
+            Union[Callable[[], Union[bytes, bytearray]], bytes, bytearray]
+        ] = None,
+    ) -> PrivateKey:
+        """Creates a PrivateKey object from a file on disk.
+
+        This method may be a convenience method that wraps ``open`` and
+        ``from_buffer``, but some TLS implementations may be able to
+        provide more-secure or faster methods of loading certificates that
+        do not involve Python code.
+
+        The ``password`` parameter behaves exactly as the equivalent
+        parameter on ``from_buffer``.
+        """
+
+
 CipherSuite = object
 DEFAULT_CIPHER_LIST = ()
 

src/mbedtls/pk.pyi

@@ -4,8 +4,10 @@
 from __future__ import annotations
 
 import enum
+import os
 import sys
 from typing import (
+    Callable,
     NamedTuple,
     NoReturn,
     Optional,
@@ -21,6 +23,12 @@ if sys.version_info < (3, 8):
 else:
     from typing import Final, Literal
 
+if sys.version_info < (3, 9):
+    _PathLike = os.PathLike
+else:
+    _PathLike = os.PathLike[str]
+
+_Path = Union[_PathLike, str]
 CIPHER_NAME: Final[Sequence[bytes]] = ...
 _DER = bytes
 _PEM = str
@@ -72,7 +80,21 @@ class CipherBase:
     def __hash__(self) -> int: ...
     def __eq__(self, other: object) -> bool: ...
     @classmethod
-    def from_buffer(cls: Type[_TCipherBase], key: bytes) -> _TCipherBase: ...
+    def from_buffer(
+        cls: Type[_TCipherBase],
+        buffer: bytes,
+        password: Optional[
+            Union[Callable[[], Union[bytes, bytearray]], bytes, bytearray]
+        ] = None,
+    ) -> _TCipherBase: ...
+    @classmethod
+    def from_file(
+        cls: Type[_TCipherBase],
+        path: _Path,
+        password: Optional[
+            Union[Callable[[], Union[bytes, bytearray]], bytes, bytearray]
+        ] = None,
+    ) -> _TCipherBase: ...
     @classmethod
     def from_DER(cls: Type[_TCipherBase], key: bytes) -> _TCipherBase: ...
     @classmethod

src/mbedtls/pk.pyx

@@ -26,6 +26,7 @@ cimport mbedtls.pk as _pk
 import enum
 from collections import namedtuple
 from functools import partial
+from pathlib import Path
 
 import mbedtls._random as _rnd
 import mbedtls.exceptions as _exc
@@ -220,7 +221,11 @@ cdef class CipherBase:
             return other.export_public_key() == self.export_public_key()
 
     @classmethod
-    def from_buffer(cls, const unsigned char[:] key not None):
+    def from_buffer(
+        cls,
+        const unsigned char[:] key not None,
+        password=None,
+    ):
         """Import a key (public or private half).
 
         The public half is generated upon importing a private key.
@@ -231,7 +236,13 @@ cdef class CipherBase:
                 password-protected private keys.
 
         """
-        raise NotImplementedError
+        if callable(password):
+            return cls(key=key, password=password())
+        return cls(key=key, password=password)
+
+    @classmethod
+    def from_file(cls, path, password=None):
+        return cls.from_buffer(Path(path).read_bytes(), password)
 
     @classmethod
     def from_DER(cls, const unsigned char[:] key not None):
@@ -531,20 +542,6 @@ cdef class RSA(CipherBase):
                  const unsigned char[:] password=None):
         super().__init__(b"RSA", key, password)
 
-    @classmethod
-    def from_buffer(cls, const unsigned char[:] key):
-        """Import a key (public or private half).
-
-        The public half is generated upon importing a private key.
-
-        Arguments:
-            key (bytes): The key in PEM or DER format.
-            password (bytes, optional): The password for
-                password-protected private keys.
-
-        """
-        return cls(key)
-
     def _has_private(self):
         """Return `True` if the key contains a valid private half."""
         return _pk.mbedtls_rsa_check_privkey(
@@ -667,20 +664,6 @@ cdef class ECC(CipherBase):
     def curve(self):
         return self._curve
 
-    @classmethod
-    def from_buffer(cls, const unsigned char[:] key not None):
-        """Import a key (public or private half).
-
-        The public half is generated upon importing a private key.
-
-        Arguments:
-            key (bytes): The key in PEM or DER format.
-            password (bytes, optional): The password for
-                password-protected private keys.
-
-        """
-        return cls(None, key)
-
     def _has_private(self):
         """Return `True` if the key contains a valid private half."""
         cdef const mbedtls_ecp_keypair* ecp = _pk.mbedtls_pk_ec(self._ctx)

src/mbedtls/tls.py

@@ -27,6 +27,7 @@
 from ._tlsi import DTLSConfiguration as DTLSConfiguration
 from ._tlsi import DTLSVersion as DTLSVersion
 from ._tlsi import NextProtocol as NextProtocol
+from ._tlsi import PrivateKey as PrivateKey
 from ._tlsi import TLSConfiguration as TLSConfiguration
 from ._tlsi import TLSVersion as TLSVersion
 
@@ -47,6 +48,7 @@
     "DTLSVersion",
     "HelloVerifyRequest",
     "NextProtocol",
+    "PrivateKey",
     "Purpose",
     "RaggedEOF",
     "ServerContext",