feat: Export import_key and generate_key methods

Author: lepture

src/joserfc/_keys.py

@@ -17,6 +17,7 @@
     "Key",
     "KeySet",
     "JWKRegistry",
+    "KeySetSerialization",
 ]
 
 Key = t.Union[OctKey, RSAKey, ECKey, OKPKey]
@@ -75,7 +76,7 @@ def import_key(cls, data: AnyKey, key_type: str | None = None, parameters: KeyPa
     def generate_key(
         cls,
         key_type: str,
-        crv_or_size: str | int,
+        crv_or_size: str | int | None = None,
         parameters: KeyParameters | None = None,
         private: bool = True,
         auto_kid: bool = False,

src/joserfc/jwk.py

@@ -5,7 +5,9 @@
     JWKRegistry,
     KeySet,
     Key,
+    KeySetSerialization,
 )
+from .rfc7517.types import AnyKey, KeyParameters
 from .rfc7518.oct_key import OctKey as OctKey
 from .rfc7518.rsa_key import RSAKey as RSAKey
 from .rfc7518.ec_key import ECKey as ECKey
@@ -19,6 +21,7 @@
     "Key",
     "KeyCallable",
     "KeyFlexible",
+    "KeySetSerialization",
     "OctKey",
     "RSAKey",
     "ECKey",
@@ -27,6 +30,8 @@
     "KeyBase",
     "GuestProtocol",
     "guess_key",
+    "import_key",
+    "generate_key",
 ]
 
 register_secp256k1()
@@ -87,3 +92,35 @@ def _normalize_key(key: KeyBase) -> Key | KeySet:
         )
         return OctKey.import_key(key)
     return key
+
+
+def import_key(
+    data: AnyKey,
+    key_type: str | None = None,
+    parameters: KeyParameters | None = None
+) -> Key:
+    """Importing a key from bytes, string, and dict. When ``value`` is a dict,
+    this method can tell the key type automatically, otherwise, developers
+    SHOULD pass the ``key_type`` themselves.
+
+    :param data: the key data in bytes, string, or dict.
+    :param key_type: an optional key type in string.
+    :param parameters: extra key parameters
+    :return: OctKey, RSAKey, ECKey, or OKPKey
+    """
+    return JWKRegistry.import_key(data, key_type, parameters)
+
+
+def generate_key(
+    key_type: str,
+    crv_or_size: str | int | None = None,
+    parameters: KeyParameters | None = None,
+    private: bool = True,
+    auto_kid: bool = False,
+) -> Key:
+    """Generating key according to the given key type. When ``key_type`` is
+    "oct" and "RSA", the second parameter SHOULD be a key size in bits.
+    When ``key_type`` is "EC" and "OKP", the second
+    parameter SHOULD be a "crv" string.
+    """
+    return JWKRegistry.generate_key(key_type, crv_or_size, parameters, private, auto_kid)

src/joserfc/rfc7518/ec_key.py

@@ -144,7 +144,11 @@ def curve_key_size(self) -> int:
 
     @classmethod
     def generate_key(
-        cls, crv: str = "P-256", parameters: KeyParameters | None = None, private: bool = True, auto_kid: bool = False
+        cls,
+        crv: str | None = "P-256",
+        parameters: KeyParameters | None = None,
+        private: bool = True,
+        auto_kid: bool = False,
     ) -> "ECKey":
         """Generate a ``ECKey`` with the given "crv" value.
 
@@ -153,6 +157,9 @@ def generate_key(
         :param private: generate a private key or public key
         :param auto_kid: add ``kid`` automatically
         """
+        if crv is None:
+            crv = "P-256"
+
         raw_key = cls.binding.generate_private_key(crv)
         if private:
             key = cls(raw_key, raw_key, parameters)

src/joserfc/rfc7518/oct_key.py

@@ -51,7 +51,11 @@ class OctKey(SymmetricKey):
 
     @classmethod
     def generate_key(
-        cls, key_size: int = 256, parameters: KeyParameters | None = None, private: bool = True, auto_kid: bool = False
+        cls,
+        key_size: int | None = 256,
+        parameters: KeyParameters | None = None,
+        private: bool = True,
+        auto_kid: bool = False,
     ) -> "OctKey":
         """Generate a ``OctKey`` with the given bit size (not bytes).
 
@@ -63,6 +67,9 @@ def generate_key(
         if not private:
             raise ValueError("oct key can not be generated as public")
 
+        if key_size is None:
+            key_size = 256
+
         if key_size % 8 != 0:
             raise ValueError("Invalid bit size for oct key")
 

src/joserfc/rfc7518/rsa_key.py

@@ -1,4 +1,5 @@
-import typing as t
+from __future__ import annotations
+from typing import TypedDict
 from functools import cached_property
 from cryptography.hazmat.primitives.asymmetric.rsa import (
     generate_private_key,
@@ -19,7 +20,7 @@
 from ..util import int_to_base64, base64_to_int
 
 
-RSADictKey = t.TypedDict(
+RSADictKey = TypedDict(
     "RSADictKey",
     {
         "n": str,
@@ -91,7 +92,7 @@ def import_public_key(obj: RSADictKey) -> RSAPublicKey:
         return numbers.public_key(default_backend())
 
     @staticmethod
-    def export_public_key(key: RSAPublicKey) -> t.Dict[str, str]:
+    def export_public_key(key: RSAPublicKey) -> dict[str, str]:
         numbers = key.public_numbers()
         return {"n": int_to_base64(numbers.n), "e": int_to_base64(numbers.e)}
 
@@ -124,16 +125,16 @@ def public_key(self) -> RSAPublicKey:
         return self.raw_value
 
     @property
-    def private_key(self) -> t.Optional[RSAPrivateKey]:
+    def private_key(self) -> RSAPrivateKey | None:
         if isinstance(self.raw_value, RSAPrivateKey):
             return self.raw_value
         return None
 
     @classmethod
     def generate_key(
         cls,
-        key_size: int = 2048,
-        parameters: t.Optional[KeyParameters] = None,
+        key_size: int | None = 2048,
+        parameters: KeyParameters | None = None,
         private: bool = True,
         auto_kid: bool = False,
     ) -> "RSAKey":
@@ -144,10 +145,15 @@ def generate_key(
         :param private: generate a private key or public key
         :param auto_kid: add ``kid`` automatically
         """
+        if key_size is None:
+            key_size = 2048
+
         if key_size < 512:
             raise ValueError("key_size must not be less than 512")
+
         if key_size % 8 != 0:
             raise ValueError("Invalid key_size for RSAKey")
+
         raw_key = generate_private_key(
             public_exponent=65537,
             key_size=key_size,

src/joserfc/rfc8037/okp_key.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import typing as t
 from functools import cached_property
 from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey, Ed25519PrivateKey
@@ -109,7 +111,7 @@ def public_key(self) -> PublicOKPKey:
         return self.raw_value
 
     @property
-    def private_key(self) -> t.Optional[PrivateOKPKey]:
+    def private_key(self) -> PrivateOKPKey | None:
         if isinstance(self.raw_value, PrivateKeyTypes):
             return self.raw_value
         return None
@@ -121,8 +123,8 @@ def curve_name(self) -> str:
     @classmethod
     def generate_key(
         cls,
-        crv: str = "Ed25519",
-        parameters: t.Optional[KeyParameters] = None,
+        crv: str | None = "Ed25519",
+        parameters: KeyParameters | None = None,
         private: bool = True,
         auto_kid: bool = False,
     ) -> "OKPKey":
@@ -133,6 +135,9 @@ def generate_key(
         :param private: generate a private key or public key
         :param auto_kid: add ``kid`` automatically
         """
+        if crv is None:
+            crv = "Ed25519"
+
         if crv not in PRIVATE_KEYS_MAP:
             raise ValueError("Invalid crv value: '{}'".format(crv))
 

tests/base.py

@@ -1,6 +1,6 @@
 import json
 import typing as t
-from joserfc.jwk import JWKRegistry
+from joserfc.jwk import import_key
 from unittest import TestCase
 from pathlib import Path
 
@@ -18,10 +18,10 @@ def load_key(filename: str, parameters=None):
 
     if filename.endswith(".json"):
         data = json.loads(content)
-        return JWKRegistry.import_key(data, parameters=parameters)
+        return import_key(data, parameters=parameters)
 
     kty = filename.split("-", 1)[0]
-    return JWKRegistry.import_key(content, kty.upper(), parameters)
+    return import_key(content, kty.upper(), parameters)
 
 
 class TestFixture(TestCase):

tests/jwk/test_key_methods.py

@@ -1,7 +1,7 @@
 from unittest import TestCase
 from joserfc.jws import register_key_set
-from joserfc.jwk import JWKRegistry, guess_key
-from joserfc.jwk import KeySet, OctKey, RSAKey
+from joserfc.jwk import guess_key, import_key, generate_key
+from joserfc.jwk import KeySet, OctKey, RSAKey, ECKey, OKPKey
 from joserfc.errors import (
     UnsupportedKeyAlgorithmError,
     UnsupportedKeyUseError,
@@ -95,24 +95,34 @@ def test_invalid_key(self):
 
     def test_import_key(self):
         # test bytes
-        key = JWKRegistry.import_key(b"secret", "oct")
+        key = import_key(b"secret", "oct")
         self.assertIsInstance(key, OctKey)
 
         # test string
-        key = JWKRegistry.import_key("secret", "oct")
+        key = import_key("secret", "oct")
         self.assertIsInstance(key, OctKey)
 
         # test dict
         data = key.as_dict()
-        key = JWKRegistry.import_key(data)
+        key = import_key(data)
         self.assertIsInstance(key, OctKey)
 
-        self.assertRaises(InvalidKeyTypeError, JWKRegistry.import_key, "secret", "invalid")
+        self.assertRaises(InvalidKeyTypeError, import_key, "secret", "invalid")
 
     def test_generate_key(self):
-        key = JWKRegistry.generate_key("oct", 8)
+        key = generate_key("oct")
         self.assertIsInstance(key, OctKey)
-        self.assertRaises(InvalidKeyTypeError, JWKRegistry.generate_key, "invalid", 8)
+
+        key = generate_key("RSA")
+        self.assertIsInstance(key, RSAKey)
+
+        key = generate_key("EC")
+        self.assertIsInstance(key, ECKey)
+
+        key = generate_key("OKP")
+        self.assertIsInstance(key, OKPKey)
+
+        self.assertRaises(InvalidKeyTypeError, generate_key, "invalid", 8)
 
     def test_check_use(self):
         key = OctKey.import_key("secret", {"use": "sig"})
@@ -141,4 +151,4 @@ def test_check_ops(self):
         self.assertRaises(UnsupportedKeyOperationError, key.check_key_op, "sign")
 
     def test_import_without_kty(self):
-        self.assertRaises(MissingKeyTypeError, JWKRegistry.import_key, {})
+        self.assertRaises(MissingKeyTypeError, import_key, {})

tests/jwk/test_oct_key.py

@@ -81,6 +81,10 @@ def test_generate_key(self):
         self.assertEqual(len(key.raw_value), 32)
         self.assertIsNone(key.kid)
 
+        key = OctKey.generate_key(None)
+        self.assertEqual(len(key.raw_value), 32)
+        self.assertIsNone(key.kid)
+
         self.assertRaises(ValueError, OctKey.generate_key, private=False)
         self.assertRaises(ValueError, OctKey.generate_key, 251)
 

tests/jwk/test_rsa_key.py

@@ -119,7 +119,7 @@ def test_generate_key(self):
         self.assertRaises(ValueError, RSAKey.generate_key, 8)
         self.assertRaises(ValueError, RSAKey.generate_key, 601)
 
-        key: RSAKey = RSAKey.generate_key(private=False)
+        key = RSAKey.generate_key(private=False)
         self.assertFalse(key.is_private)
         self.assertIsNone(key.kid)