feat(jwk): add a derive_key method for OKPKey

#82

Author: lepture

src/joserfc/_rfc8037/okp_key.py

@@ -2,6 +2,8 @@
 
 importtypingast
 fromfunctoolsimportcached_property
+fromcryptography.hazmat.primitivesimporthashes
+fromcryptography.hazmat.backendsimportdefault_backend
 fromcryptography.hazmat.primitives.asymmetric.ed25519importEd25519PublicKey, Ed25519PrivateKey
 fromcryptography.hazmat.primitives.asymmetric.ed448importEd448PublicKey, Ed448PrivateKey
 fromcryptography.hazmat.primitives.asymmetric.x25519importX25519PublicKey, X25519PrivateKey
@@ -12,6 +14,8 @@
 PrivateFormat,
 NoEncryption,
 )
+fromcryptography.hazmat.primitives.kdf.hkdfimportHKDF
+fromcryptography.hazmat.primitives.kdf.pbkdf2importPBKDF2HMAC
 from .._rfc7517.modelsimportCurveKey
 from .._rfc7517.typesimportKeyParameters, AnyKey
 from .._rfc7517.pemimportCryptographyBinding
@@ -44,6 +48,12 @@
 "X25519": X25519PrivateKey,
 "X448": X448PrivateKey,
 }
+OKP_SEED_SIZES: dict[LiteralCurves, int] = {
+"Ed25519": 32,
+"Ed448": 57,
+"X25519": 32,
+"X448": 56,
+}
 PrivateKeyTypes= (Ed25519PrivateKey, Ed448PrivateKey, X25519PrivateKey, X448PrivateKey)
 
 
@@ -193,14 +203,103 @@ def generate_key(
 raw_key=cls.binding.generate_private_key("Ed25519")
 else:
 raw_key=cls.binding.generate_private_key(crv)
-ifprivate:
-key=cls(raw_key, raw_key, parameters)
+return_wrap_key(cls, raw_key, private, auto_kid, parameters)
+
+@classmethod
+defderive_key(
+cls: t.Type["OKPKey"],
+secret: bytes|str,
+crv: LiteralCurves="Ed25519",
+parameters: KeyParameters|None=None,
+private: bool=True,
+auto_kid: bool=False,
+kdf_name: t.Literal["HKDF", "PBKDF2"] ="HKDF",
+kdf_options: dict[str, t.Any] |None=None,
+ ) ->"OKPKey":
+"""
+ Derives a key from a given input secret using a specified key derivation function
+ (KDF) and elliptic curve algorithm.
+
+ To derive a key using **HKDF**, the ``kdf_options`` may contain the ``algorithm``,
+ ``salt`` and ``info`` values:
+
+ .. code-block:: python
+
+ from cryptography.hazmat.primitives import hashes
+ from joserfc.jwk import OKPKey
+
+ # default kdf_name is HKDF, algorithm is SHA256
+ OKPKey.derive_key("secret")
+ # equivalent to
+ OKPKey.derive_key(
+ "secret", "Ed25519",
+ kdf_name="HKDF",
+ kdf_options={
+ "algorithm": hashes.SHA256(),
+ "salt": b"joserfc:OKP:HKDF:Ed25519",
+ "info": b"",
+ }
+ )
+
+ To derive a key using **PBKDF2**, the ``kdf_options`` may contain the ``algorithm``,
+ ``salt`` and ``iterations`` values:
+
+ .. code-block:: python
+
+ from cryptography.hazmat.primitives import hashes
+ from joserfc.jwk import OKPKey
+
+ OKPKey.derive_key("secret", kdf_name="PBKDF2")
+ # equivalent to
+ OKPKey.derive_key(
+ "secret", "Ed25519",
+ kdf_name="PBKDF2",
+ kdf_options={
+ "algorithm": hashes.SHA256(),
+ "salt": b"joserfc:OKP:PBKDF2:Ed25519",
+ "iterations": 100000,
+ }
+ )
+
+ :param secret: The input secret used for key derivation.
+ :param crv: OKPKey curve name
+ :param parameters: extra parameter in JWK
+ :param private: generate a private key or public key
+ :param auto_kid: add ``kid`` automatically
+ :param kdf_name: Key derivation function name
+ :param kdf_options: Additional options for the KDF
+ """
+ifkdf_optionsisNone:
+kdf_options= {}
+
+algorithm=kdf_options.pop("algorithm", None)
+ifalgorithmisNone:
+algorithm=hashes.SHA256()
+
+kdf_options.setdefault("salt", to_bytes(f"joserfc:OKP:{kdf_name}:{crv}"))
+ifkdf_name=="HKDF":
+kdf_options.setdefault("info", b"")
+hkdf=HKDF(
+algorithm=algorithm,
+length=OKP_SEED_SIZES[crv],
+backend=default_backend(),
+**kdf_options,
+ )
+seed=hkdf.derive(to_bytes(secret))
+elifkdf_name=="PBKDF2":
+kdf_options.setdefault("iterations", 100000)
+pbkdf2=PBKDF2HMAC(
+algorithm=algorithm,
+length=OKP_SEED_SIZES[crv],
+backend=default_backend(),
+**kdf_options,
+ )
+seed=pbkdf2.derive(to_bytes(secret))
 else:
-pub_key=raw_key.public_key()
-key=cls(pub_key, pub_key, parameters)
-ifauto_kid:
-key.ensure_kid()
-returnkey
+raiseValueError(f"Invalid kdf value: '{kdf_name}'")
+
+raw_key=cls.binding.from_private_bytes(crv, seed)
+return_wrap_key(cls, raw_key, private, auto_kid, parameters)
 
 
 defget_key_curve(key: t.Union[PublicOKPKey, PrivateOKPKey]) ->LiteralCurves:
@@ -213,3 +312,20 @@ def get_key_curve(key: t.Union[PublicOKPKey, PrivateOKPKey]) -> LiteralCurves:
 elifisinstance(key, (X448PublicKey, X448PrivateKey)):
 return"X448"
 raiseValueError("Invalid key") # pragma: no cover
+
+
+def_wrap_key(
+cls: t.Type["OKPKey"],
+raw_key: PrivateOKPKey,
+private: bool,
+auto_kid: bool,
+parameters: KeyParameters|None=None,
+) ->OKPKey:
+ifprivate:
+key=cls(raw_key, raw_key, parameters)
+else:
+pub_key=raw_key.public_key()
+key=cls(pub_key, pub_key, parameters)
+ifauto_kid:
+key.ensure_kid()
+returnkey

tests/jwk/test_okp_key.py

@@ -1,4 +1,5 @@
 fromunittestimportTestCase
+fromcryptography.hazmat.primitivesimporthashes
 fromjoserfc.jwkimportOKPKey
 fromjoserfc.errorsimport (
 InvalidExchangeKeyError,
@@ -118,3 +119,31 @@ def test_key_eq(self):
 self.assertEqual(key1, key2)
 key3=OKPKey.generate_key()
 self.assertNotEqual(key1, key3)
+
+deftest_derive_key_errors(self):
+self.assertRaises(KeyError, OKPKey.derive_key, "secret", "invalid")
+self.assertRaises(ValueError, OKPKey.derive_key, "secret", "Ed25519", kdf_name="invalid")
+
+deftest_derive_key_with_default_kwargs(self):
+curves= ["Ed25519", "Ed448", "X25519", "X448"]
+forcrvincurves:
+key1=OKPKey.derive_key("okp-secret-key", crv)
+key2=OKPKey.derive_key("okp-secret-key", crv)
+self.assertEqual(key1, key2)
+
+forcrvincurves:
+key1=OKPKey.derive_key("okp-secret-key", crv, kdf_name="PBKDF2")
+key2=OKPKey.derive_key("okp-secret-key", crv, kdf_name="PBKDF2")
+self.assertEqual(key1, key2)
+
+deftest_derive_key_with_new_salt(self):
+curves= ["Ed25519", "Ed448", "X25519", "X448"]
+forcrvincurves:
+key1=OKPKey.derive_key("okp-secret-key", crv, kdf_options={"salt": b"salt"})
+key2=OKPKey.derive_key("okp-secret-key", crv, kdf_options={"salt": b"salt"})
+self.assertEqual(key1, key2)
+
+deftest_derive_key_with_different_hash(self):
+key1=OKPKey.derive_key("okp-secret-key", "Ed25519", kdf_options={"algorithm": hashes.SHA256()})
+key2=OKPKey.derive_key("okp-secret-key", "Ed25519", kdf_options={"algorithm": hashes.SHA512()})
+self.assertNotEqual(key1, key2)