feat: import bytes/str keys without specified key type, #73

Author: lepture

src/joserfc/_rfc7517/pem.py

@@ -23,7 +23,11 @@
 from ..util import to_bytes
 
 
-def load_pem_key(raw: bytes, password: bytes | None = None) -> Any:
+def import_from_ssh_key(raw: bytes) -> Any:
+    return load_ssh_public_key(raw, backend=default_backend())
+
+
+def import_from_pem_key(raw: bytes, password: bytes | None = None) -> Any:
     key: Any
 
     if b"OPENSSH PRIVATE" in raw:
@@ -91,15 +95,15 @@ def dump_pem_key(
 class CryptographyBinding(NativeKeyBinding, metaclass=ABCMeta):
     key_type: str
     ssh_type: bytes
-    cryptography_native_keys: Any
+    _cryptography_key_types: Any
 
     @classmethod
     def check_ssh_type(cls, value: bytes) -> bool:
         return value.startswith(cls.ssh_type)
 
     @classmethod
-    def check_cryptography_native_key(cls, native_key: Any) -> bool:
-        return isinstance(native_key, cls.cryptography_native_keys)
+    def check_cryptography_key(cls, native_key: Any) -> bool:
+        return isinstance(native_key, cls._cryptography_key_types)
 
     @classmethod
     def convert_raw_key_to_dict(cls, raw_key: Any, private: bool) -> DictKey:
@@ -118,13 +122,13 @@ 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())
+            return import_from_ssh_key(value)
 
         if password is not None:
             password = to_bytes(password)
 
-        key = load_pem_key(value, password)
-        if not cls.check_cryptography_native_key(key):
+        key = import_from_pem_key(value, password)
+        if not cls.check_cryptography_key(key):
             raise InvalidKeyTypeError(f"Not a key of: '{cls.key_type}'")
         return key
 
@@ -135,11 +139,13 @@ def as_bytes(
         private: bool | None = False,
         password: Any | None = None,
     ) -> bytes:
+        if private is None:
+            private = key.is_private
+
         if private:
             return dump_pem_key(key.private_key, encoding, private, password)
-        elif private is False:
+        else:
             return dump_pem_key(key.public_key, encoding, private, password)
-        return dump_pem_key(key.raw_value, encoding, key.is_private, password)
 
     @staticmethod
     @abstractmethod

src/joserfc/_rfc7518/ec_key.py

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

src/joserfc/_rfc7518/rsa_key.py

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

src/joserfc/_rfc8037/okp_key.py

@@ -48,7 +48,7 @@
 class OKPBinding(CryptographyBinding):
     key_type = "OKP"
     ssh_type = b"ssh-ed25519"
-    cryptography_native_keys = (
+    _cryptography_key_types = (
         Ed25519PublicKey,
         Ed25519PrivateKey,
         Ed448PublicKey,

src/joserfc/jwk.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 import typing as t
+import warnings
 from ._keys import (
     JWKRegistry,
     KeySet,
     Key,
     KeySetSerialization,
 )
+from ._rfc7517.pem import import_from_pem_key, import_from_ssh_key
 from ._rfc7517.types import AnyKey, DictKey, KeyParameters
 from ._rfc7518.oct_key import OctKey
 from ._rfc7518.rsa_key import RSAKey
@@ -14,7 +16,9 @@
 from ._rfc8812 import register_secp256k1
 from ._rfc7638 import calculate_thumbprint as thumbprint
 from ._rfc9278 import calculate_thumbprint_uri as thumbprint_uri
+from .errors import SecurityWarning, InvalidKeyTypeError
 from .registry import Header
+from .util import to_bytes
 
 
 __all__ = [
@@ -116,7 +120,7 @@ def import_key(data: AnyKey, key_type: t.Literal["OKP"], parameters: KeyParamete
 
 
 @t.overload
-def import_key(data: DictKey, key_type: None = None, parameters: KeyParameters | None = None) -> Key: ...
+def import_key(data: AnyKey, key_type: None = None, parameters: KeyParameters | None = None) -> Key: ...
 
 
 def import_key(
@@ -133,6 +137,28 @@ def import_key(
     :param parameters: extra key parameters
     :return: OctKey, RSAKey, ECKey, or OKPKey
     """
+    if isinstance(data, (str, bytes)) and key_type is None:
+        warnings.warn("Using implicit key type is not recommended.", SecurityWarning)
+
+        value = to_bytes(data)
+        ssh_types = tuple(
+            cls.binding.ssh_type for cls in JWKRegistry.key_types.values() if hasattr(cls.binding, "ssh_type")
+        )
+        if value.startswith(ssh_types):
+            try:
+                raw_key = import_from_ssh_key(value)
+            except ValueError:
+                return OctKey.import_key(value, parameters)
+        else:
+            try:
+                raw_key = import_from_pem_key(value)
+            except ValueError:
+                return OctKey.import_key(value, parameters)
+
+        for cls in JWKRegistry.key_types.values():
+            if hasattr(cls.binding, "check_cryptography_key") and cls.binding.check_cryptography_key(raw_key):
+                return cls(raw_key, data, parameters)
+        raise InvalidKeyTypeError("Not a key of any supported type")  # pragma: no cover
     return JWKRegistry.import_key(data, key_type, parameters)
 
 

tests/jwk/test_key_methods.py

@@ -1,7 +1,12 @@
 from unittest import TestCase
 
 from joserfc import jws
-from joserfc.jwk import guess_key, import_key, generate_key, thumbprint_uri
+from joserfc.jwk import (
+    guess_key,
+    import_key,
+    generate_key,
+    thumbprint_uri,
+)
 from joserfc.jwk import KeySet, OctKey, RSAKey, ECKey, OKPKey
 from joserfc.errors import (
     UnsupportedKeyAlgorithmError,
@@ -11,6 +16,7 @@
     MissingKeyTypeError,
     InvalidKeyIdError,
 )
+from tests.keys import read_key
 
 
 class Guest:
@@ -170,3 +176,25 @@ def test_find_correct_key_with_alg(self):
         jws.serialize_compact({"alg": "HS256"}, "foo", key_set)
         # get by kid
         jws.serialize_compact({"alg": "HS256", "kid": key2.kid}, "foo", key_set)
+
+    def test_import_bytes_keys(self):
+        # check by ssh types
+        key = import_key(read_key("ssh-rsa-public.pem"))
+        self.assertEqual(key.key_type, "RSA")
+        key = import_key(read_key("ssh-ecdsa-public.pem"))
+        self.assertEqual(key.key_type, "EC")
+        key = import_key(read_key("ssh-ed25519-public.pem"))
+        self.assertEqual(key.key_type, "OKP")
+        key = import_key("ssh-rsa-oct")
+        self.assertEqual(key.key_type, "oct")
+
+        # check by pem types
+        key = import_key(read_key("rsa-openssl-public.pem"))
+        self.assertEqual(key.key_type, "RSA")
+        key = import_key(read_key("ec-p256-public.pem"))
+        self.assertEqual(key.key_type, "EC")
+        key = import_key(read_key("okp-ed448-public.pem"))
+        self.assertEqual(key.key_type, "OKP")
+
+        key = import_key(b"oct-key")
+        self.assertEqual(key.key_type, "oct")

tests/keys/ssh-ecdsa-public.pem

@@ -0,0 +1 @@
+ecdsa-sha2-nistp256 AAAAE2VjZHNhLXNoYTItbmlzdHAyNTYAAAAIbmlzdHAyNTYAAABBBLTEVnhsZrRnFhOlXN68wUslJD6/Xo5N4TprdkyiG3UZYKN7tT8ot0ZLP7XOv3lusPd9+A128hRjYhKQIFHPCyY=

tests/keys/ssh-ed25519-public.pem

@@ -0,0 +1 @@
+ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAILuDaT2DEQNFTGvcnOsDp+cvSppjVRPajco10eDh76LB