fix(jwk): allow import key from cryptography native key types

Author: lepture

src/joserfc/_rfc7518/ec_key.py

@@ -147,15 +147,20 @@ def curve_key_size(self) -> int:
     @classmethod
     def import_key(
         cls: t.Any,
-        value: AnyKey,
+        value: AnyKey | EllipticCurvePrivateKey | EllipticCurvePublicKey,
         parameters: KeyParameters | None = None,
         password: t.Any = None,
     ) -> "ECKey":
-        return super(ECKey, cls).import_key(value, parameters, password)
+        key: ECKey
+        if isinstance(value, (EllipticCurvePrivateKey, EllipticCurvePublicKey)):
+            key = cls(value, value, parameters)
+        else:
+            key = super(ECKey, cls).import_key(value, parameters, password)
+        return key
 
     @classmethod
     def generate_key(
-        cls,
+        cls: t.Type["ECKey"],
         crv: str | None = "P-256",
         parameters: KeyParameters | None = None,
         private: bool = True,

src/joserfc/_rfc7518/rsa_key.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import warnings
-from typing import TypedDict, Any
+import typing as t
 from functools import cached_property
 from cryptography.hazmat.primitives.asymmetric.rsa import (
     generate_private_key,
@@ -22,7 +22,7 @@
 from ..util import int_to_base64, base64_to_int
 
 
-RSADictKey = TypedDict(
+RSADictKey = t.TypedDict(
     "RSADictKey",
     {
         "n": str,
@@ -136,20 +136,24 @@ def private_key(self) -> RSAPrivateKey | None:
 
     @classmethod
     def import_key(
-        cls: Any,
-        value: AnyKey,
+        cls: t.Any,
+        value: AnyKey | RSAPrivateKey | RSAPublicKey,
         parameters: KeyParameters | None = None,
-        password: Any = None,
+        password: t.Any = None,
     ) -> "RSAKey":
-        key: RSAKey = super(RSAKey, cls).import_key(value, parameters, password)
+        key: RSAKey
+        if isinstance(value, (RSAPrivateKey, RSAPublicKey)):
+            key = cls(value, value, parameters)
+        else:
+            key = super(RSAKey, cls).import_key(value, parameters, password)
         if key.raw_value.key_size < 2048:
             # https://csrc.nist.gov/publications/detail/sp/800-131a/rev-2/final
             warnings.warn("Key size should be >= 2048 bits", SecurityWarning)
         return key
 
     @classmethod
     def generate_key(
-        cls,
+        cls: t.Type["RSAKey"],
         key_size: int | None = 2048,
         parameters: KeyParameters | None = None,
         private: bool = True,

src/joserfc/_rfc8037/okp_key.py

@@ -150,11 +150,12 @@ def curve_name(self) -> LiteralCurves:
 
     @classmethod
     def import_key(
-        cls: t.Type[OKPKey],
+        cls: t.Any,
         value: AnyKey | PrivateOKPKey | PublicOKPKey,
         parameters: KeyParameters | None = None,
         password: t.Any = None,
     ) -> "OKPKey":
+        key: OKPKey
         if isinstance(
             value,
             (
@@ -168,12 +169,14 @@ def import_key(
                 X448PublicKey,
             ),
         ):
-            return cls(value, value, parameters)
-        return super(OKPKey, cls).import_key(value, parameters, password)
+            key = cls(value, value, parameters)
+        else:
+            key = super(OKPKey, cls).import_key(value, parameters, password)
+        return key
 
     @classmethod
     def generate_key(
-        cls: t.Type[OKPKey],
+        cls: t.Type["OKPKey"],
         crv: LiteralCurves | None = "Ed25519",
         parameters: KeyParameters | None = None,
         private: bool = True,

tests/jwk/test_ec_key.py

@@ -45,6 +45,10 @@ def test_import_p512_key(self):
     def test_import_secp256k1_key(self):
         self.run_import_key("secp256k1")
 
+    def test_import_from_native_keys(self):
+        key = ECKey.generate_key()
+        self.assertEqual(key, ECKey.import_key(key.private_key))
+
     def test_generate_key(self):
         self.assertRaises(ValueError, ECKey.generate_key, "Invalid")
 

tests/jwk/test_rsa_key.py

@@ -92,6 +92,10 @@ def test_import_key_from_openssl(self):
         key: RSAKey = RSAKey.import_key(private_pem)
         self.assertTrue(key.is_private)
 
+    def test_import_from_native_keys(self):
+        key = RSAKey.generate_key()
+        self.assertEqual(key, RSAKey.import_key(key.private_key))
+
     def test_output_as_methods(self):
         private_pem = read_key("rsa-openssl-private.pem")
         key: RSAKey = RSAKey.import_key(private_pem)