add functions to generate a similar key (same kty, crv, size)

latchset · Dec 2, 2024 · 3fb38f2 · 3fb38f2
1 parent f105494
commit 3fb38f2
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 0 deletions.
diff --git a/jwcrypto/jwk.py b/jwcrypto/jwk.py
@@ -368,6 +368,18 @@ def generate_key(self, **params):
 
         gen(params)
 
+    @classmethod
+    def generate_similar(cls, key):
+         return cls.generate(**key.get_generate_params())
+
+    def get_generate_params(self):
+         params = {param: self.get(param) for param in ["kty", "crv"] if param in self}
+         if self.get("kty") == "RSA":
+             params["size"] = self._get_public_key().key_size
+         elif self.get("kty") == "oct":
+             params["size"] = len(base64url_decode(self.k)) * 8
+         return params
+
     def _get_gen_size(self, params, default_size=None):
         size = default_size
         if 'size' in params:

diff --git a/jwcrypto/tests.py b/jwcrypto/tests.py
@@ -506,6 +506,48 @@ def test_jwk_from_json(self):
         y = jwk.JWK.from_json(k.export())
         self.assertEqual(k.export(), y.export())
 
+    def test_generate_similar(self):
+        KEY_PARAMS = [
+            {
+                "kty": "oct",
+                "size": 192
+            },
+            {
+                "kty": "RSA",
+                "size": 3072
+            },
+            {
+                "kty": "EC",
+                "crv": "P-256"
+            },
+            {
+                "kty": "EC",
+                "crv": "P-384",
+            },
+            {
+                "kty": "OKP",
+                "crv": "Ed25519"
+            },
+            {
+                "kty": "OKP",
+                "crv": "Ed448"
+            },
+            {
+                "kty": "OKP",
+                "crv": "X25519"
+            },
+            {
+                "kty": "OKP",
+                "crv": "X448"
+            }
+        ]
+        for params in KEY_PARAMS:
+            key1 = jwk.JWK.generate(**params)
+            key2 = jwk.JWK.generate_similar(key1)
+            self.assertEqual(key1.get("kty"), key2.get("kty"))
+            self.assertEqual(key1.get("crv"), key2.get("crv"))
+            self.assertEqual(type(key1.get_op_key("sign")), type(key2.get_op_key("sign")))
+
     def test_jwkset(self):
         k = jwk.JWK(**RSAPrivateKey)
         ks = jwk.JWKSet()