add key wrapping functions

Author: jschlyter

src/cryptojwt/jwk/wrap.py

@@ -0,0 +1,31 @@
+"""JWK wrapping"""
+
+import json
+
+from . import JWK
+from ..jwe.jwe import JWE
+from ..jwx import key_from_jwk_dict
+
+__author__ = 'jschlyter'
+
+DEFAULT_WRAP_PARAMS = {
+    "EC": {"alg": "ECDH-ES+A128KW", "enc": "A128GCM"},
+    "RSA": {"alg": "RSA1_5", "enc": "A128CBC-HS256"},
+    "oct": {"alg": "A128KW", "enc": "A128CBC-HS256"},
+}
+
+
+def wrap_key(key: JWK, wrapping_key: JWK, wrap_params: dict = DEFAULT_WRAP_PARAMS) -> str:
+    message = json.dumps(key.serialize(private=True)).encode()
+    try:
+        enc_params = wrap_params[wrapping_key.kty]
+    except KeyError:
+        raise ValueError("Unsupported wrapping key type")
+    _jwe = JWE(msg=message, **enc_params)
+    return _jwe.encrypt(keys=[wrapping_key], kid=wrapping_key.kid)
+
+
+def unwrap_key(jwe: str, wrapping_keys: JWK) -> JWK:
+    _jwe = JWE()
+    message = _jwe.decrypt(token=jwe, keys=wrapping_keys)
+    return key_from_jwk_dict(json.loads(message))

tests/test_10_jwk_wrap.py

@@ -0,0 +1,42 @@
+import os
+
+from cryptojwt.jwk.ec import new_ec_key
+from cryptojwt.jwk.hmac import SYMKey
+from cryptojwt.jwk.rsa import new_rsa_key
+from cryptojwt.jwk.wrap import wrap_key, unwrap_key
+
+__author__ = 'jschlyter'
+
+WRAPPING_KEYS = [
+    SYMKey(use="enc", key=os.urandom(32)),
+    new_ec_key(crv="P-256"),
+    new_ec_key(crv="P-384"),
+    new_rsa_key(size=2048),
+    new_rsa_key(size=4096),
+]
+
+SECRET_KEYS = [
+    SYMKey(use="enc", key=os.urandom(32)),
+    new_ec_key(crv="P-256"),
+    new_rsa_key(size=2048),
+]
+
+
+def test_wrap_default():
+    for wrapping_key in WRAPPING_KEYS:
+        for key in SECRET_KEYS:
+            wrapped_key = wrap_key(key, wrapping_key)
+            unwrapped_key = unwrap_key(wrapped_key, [wrapping_key])
+            assert key == unwrapped_key
+
+def test_wrap_params():
+    wrap_params = {
+        "EC": {"alg": "ECDH-ES+A256KW", "enc": "A256GCM"},
+        "RSA": {"alg": "RSA1_5", "enc": "A256CBC-HS512"},
+        "oct": {"alg": "A256KW", "enc": "A256CBC-HS512"},
+    }
+    for wrapping_key in WRAPPING_KEYS:
+        for key in SECRET_KEYS:
+            wrapped_key = wrap_key(key, wrapping_key, wrap_params)
+            unwrapped_key = unwrap_key(wrapped_key, [wrapping_key])
+            assert key == unwrapped_key