feat: check ssh_type and cryptography key types in bindings

Author: lepture

docs/migrations/authlib.rst

@@ -80,7 +80,6 @@ When using methods such as ``.as_dict``, ``.as_bytes``, ``.as_pem``, and others,
 
 .. code-block:: python
     :caption: Authlib
-    :emphasize-lines: 1,2
 
     key.as_dict(is_private=True)
 

src/joserfc/_rfc7517/pem.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import Any, Literal, cast
+from typing import Any, Literal, Tuple, cast
 from abc import ABCMeta, abstractmethod
 from cryptography.x509 import load_pem_x509_certificate
 from cryptography.hazmat.primitives.serialization import (
@@ -19,15 +19,14 @@
 from cryptography.hazmat.backends import default_backend
 from .models import NativeKeyBinding, GenericKey
 from .types import DictKey
+from ..errors import InvalidKeyTypeError
 from ..util import to_bytes
 
 
-def load_pem_key(raw: bytes, ssh_type: bytes | None = None, password: bytes | None = None) -> Any:
+def load_pem_key(raw: bytes, password: bytes | None = None) -> Any:
     key: Any
-    if ssh_type and raw.startswith(ssh_type):
-        key = load_ssh_public_key(raw, backend=default_backend())
 
-    elif b"OPENSSH PRIVATE" in raw:
+    if b"OPENSSH PRIVATE" in raw:
         key = load_ssh_private_key(raw, password=password, backend=default_backend())
 
     elif b"PUBLIC" in raw:
@@ -49,7 +48,10 @@ def load_pem_key(raw: bytes, ssh_type: bytes | None = None, password: bytes | No
 
 
 def dump_pem_key(
-    key: Any, encoding: Literal["PEM", "DER"] | None = None, private: bool | None = False, password: Any | None = None
+    key: Any,
+    encoding: Literal["PEM", "DER"] | None = None,
+    private: bool | None = False,
+    password: Any | None = None,
 ) -> bytes:
     """Export key into PEM/DER format bytes.
 
@@ -87,7 +89,17 @@ def dump_pem_key(
 
 
 class CryptographyBinding(NativeKeyBinding, metaclass=ABCMeta):
+    key_type: str
     ssh_type: bytes
+    cryptography_native_keys: Tuple[Any]
+
+    @classmethod
+    def check_ssh_type(cls, value: bytes):
+        return cls.ssh_type and value.startswith(cls.ssh_type)
+
+    @classmethod
+    def check_cryptography_native_key(cls, native_key: Any):
+        return isinstance(native_key, cls.cryptography_native_keys)
 
     @classmethod
     def convert_raw_key_to_dict(cls, raw_key: Any, private: bool) -> DictKey:
@@ -105,9 +117,16 @@ def import_from_dict(cls, value: DictKey) -> Any:
 
     @classmethod
     def import_from_bytes(cls, value: bytes, password: Any | None = None) -> Any:
+        if cls.check_ssh_type(value):
+            return load_ssh_public_key(value, backend=default_backend())
+
         if password is not None:
             password = to_bytes(password)
-        return load_pem_key(value, cls.ssh_type, password)
+
+        key = load_pem_key(value, password)
+        if not cls.check_cryptography_native_key(key):
+            raise InvalidKeyTypeError(f"Not a key of: '{cls.key_type}'")
+        return key
 
     @staticmethod
     def as_bytes(
@@ -116,7 +135,7 @@ def as_bytes(
         private: bool | None = False,
         password: Any | None = None,
     ) -> bytes:
-        if private is True:
+        if private:
             return dump_pem_key(key.private_key, encoding, private, password)
         elif private is False:
             return dump_pem_key(key.public_key, encoding, private, password)

src/joserfc/_rfc7518/ec_key.py

@@ -36,7 +36,9 @@
 
 
 class ECBinding(CryptographyBinding):
+    key_type = "EC"
     ssh_type = b"ecdsa-sha2-"
+    cryptography_native_keys = (EllipticCurvePrivateKey, EllipticCurvePublicKey)
 
     _dss_curves: dict[str, t.Type[EllipticCurve]] = {}
     _curves_dss: dict[str, str] = {}

src/joserfc/_rfc7518/rsa_key.py

@@ -39,7 +39,9 @@
 
 
 class RSABinding(CryptographyBinding):
+    key_type = "RSA"
     ssh_type = b"ssh-rsa"
+    cryptography_native_keys = (RSAPrivateKey, RSAPublicKey)
 
     @staticmethod
     def import_private_key(obj: RSADictKey) -> RSAPrivateKey:

src/joserfc/_rfc8037/okp_key.py

@@ -46,7 +46,18 @@
 
 
 class OKPBinding(CryptographyBinding):
+    key_type = "OKP"
     ssh_type = b"ssh-ed25519"
+    cryptography_native_keys = (
+        Ed25519PublicKey,
+        Ed25519PrivateKey,
+        Ed448PublicKey,
+        Ed448PrivateKey,
+        X25519PublicKey,
+        X25519PrivateKey,
+        X448PublicKey,
+        X448PrivateKey,
+    )
 
     @staticmethod
     def import_private_key(obj: OKPDictKey) -> PrivateOKPKey: