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__ importannotations
-fromtypingimportAny, Literal, cast
+fromtypingimportAny, Literal, Tuple, cast
 fromabcimportABCMeta, abstractmethod
 fromcryptography.x509importload_pem_x509_certificate
 fromcryptography.hazmat.primitives.serializationimport (
@@ -19,15 +19,14 @@
 fromcryptography.hazmat.backendsimportdefault_backend
 from .modelsimportNativeKeyBinding, GenericKey
 from .typesimportDictKey
+from ..errorsimportInvalidKeyTypeError
 from ..utilimportto_bytes
 
 
-defload_pem_key(raw: bytes, ssh_type: bytes|None=None, password: bytes|None=None) ->Any:
+defload_pem_key(raw: bytes, password: bytes|None=None) ->Any:
 key: Any
-ifssh_typeandraw.startswith(ssh_type):
-key=load_ssh_public_key(raw, backend=default_backend())
 
-elifb"OPENSSH PRIVATE"inraw:
+ifb"OPENSSH PRIVATE"inraw:
 key=load_ssh_private_key(raw, password=password, backend=default_backend())
 
 elifb"PUBLIC"inraw:
@@ -49,7 +48,10 @@ def load_pem_key(raw: bytes, ssh_type: bytes | None = None, password: bytes | No
 
 
 defdump_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(
 
 
 classCryptographyBinding(NativeKeyBinding, metaclass=ABCMeta):
+key_type: str
 ssh_type: bytes
+cryptography_native_keys: Tuple[Any]
+
+@classmethod
+defcheck_ssh_type(cls, value: bytes):
+returncls.ssh_typeandvalue.startswith(cls.ssh_type)
+
+@classmethod
+defcheck_cryptography_native_key(cls, native_key: Any):
+returnisinstance(native_key, cls.cryptography_native_keys)
 
 @classmethod
 defconvert_raw_key_to_dict(cls, raw_key: Any, private: bool) ->DictKey:
@@ -105,9 +117,16 @@ def import_from_dict(cls, value: DictKey) -> Any:
 
 @classmethod
 defimport_from_bytes(cls, value: bytes, password: Any|None=None) ->Any:
+ifcls.check_ssh_type(value):
+returnload_ssh_public_key(value, backend=default_backend())
+
 ifpasswordisnotNone:
 password=to_bytes(password)
-returnload_pem_key(value, cls.ssh_type, password)
+
+key=load_pem_key(value, password)
+ifnotcls.check_cryptography_native_key(key):
+raiseInvalidKeyTypeError(f"Not a key of: '{cls.key_type}'")
+returnkey
 
 @staticmethod
 defas_bytes(
@@ -116,7 +135,7 @@ def as_bytes(
 private: bool|None=False,
 password: Any|None=None,
  ) ->bytes:
-ifprivateisTrue:
+ifprivate:
 returndump_pem_key(key.private_key, encoding, private, password)
 elifprivateisFalse:
 returndump_pem_key(key.public_key, encoding, private, password)

src/joserfc/_rfc7518/ec_key.py

@@ -36,7 +36,9 @@
 
 
 classECBinding(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 @@
 
 
 classRSABinding(CryptographyBinding):
+key_type="RSA"
 ssh_type=b"ssh-rsa"
+cryptography_native_keys= (RSAPrivateKey, RSAPublicKey)
 
 @staticmethod
 defimport_private_key(obj: RSADictKey) ->RSAPrivateKey:

src/joserfc/_rfc8037/okp_key.py

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