Stricter

Author: rohe

src/cryptojwt/jwe/jwe.py

@@ -1,7 +1,10 @@
 import logging
 
-from cryptojwt.jwk.jwk import key_from_jwk_dict
 from ..jwk.asym import AsymmetricKey
+from ..jwk.ec import ECKey
+from ..jwk.hmac import SYMKey
+from ..jwk.jwk import key_from_jwk_dict
+from ..jwk.rsa import RSAKey
 from ..jwx import JWx
 
 from .exception import DecryptionFailed
@@ -82,14 +85,16 @@ def encrypt(self, keys=None, cek="", iv="", **kwargs):
 
         # Determine Encryption Class by Algorithm
         if _alg in ["RSA-OAEP", "RSA-OAEP-256", "RSA1_5"]:
+            keys = [k for k in keys if isinstance(k, RSAKey)]
             encrypter = JWE_RSA(self.msg, **self._dict)
         elif _alg.startswith("A") and _alg.endswith("KW"):
+            keys = [k for k in keys if isinstance(k, SYMKey)]
             encrypter = JWE_SYM(self.msg, **self._dict)
         elif _alg.startswith("ECDH-ES"):
-
-            # ECDH-ES Requires the Server ECDH-ES Key to be set
+            keys = [k for k in keys if isinstance(k, ECKey)]
             if not keys:
-                raise NoSuitableECDHKey(_alg)
+                logger.error(KEY_ERR.format(_alg))
+                raise NoSuitableEncryptionKey(_alg)
 
             encrypter = JWE_EC(**self._dict)
             cek, encrypted_key, iv, params, eprivk = encrypter.enc_setup(
@@ -100,17 +105,25 @@ def encrypt(self, keys=None, cek="", iv="", **kwargs):
             logger.error("'{}' is not a supported algorithm".format(_alg))
             raise NotSupportedAlgorithm
 
+        if not keys:
+            logger.error(KEY_ERR.format(_alg))
+            raise NoSuitableEncryptionKey(_alg)
+
         if cek:
             kwargs["cek"] = cek
 
         if iv:
             kwargs["iv"] = iv
 
         for key in keys:
-            if isinstance(key, AsymmetricKey):
+            if isinstance(key, SYMKey):
+                _key = key.key
+            elif isinstance(key, ECKey):
                 _key = key.public_key()
+            elif isinstance(key, RSAKey):
+                    _key = key.public_key()
             else:
-                _key = key.key
+                raise ValueError('Unknown key type')
 
             if key.kid:
                 encrypter["kid"] = key.kid

src/cryptojwt/jwe/jwe_hmac.py

@@ -11,7 +11,7 @@
 from ..exception import MissingKey
 from ..exception import WrongNumberOfParts
 from ..jwk.hmac import SYMKey
-from ..utils import intarr2str
+from ..utils import intarr2str, as_bytes
 
 logger = logging.getLogger(__name__)
 
@@ -31,7 +31,7 @@ def encrypt(self, key, iv="", cek="", **kwargs):
         :param kwargs: Extra keyword arguments, just ignore for now.
         :return:
         """
-        _msg = self.msg
+        _msg = as_bytes(self.msg)
 
         _args = self._dict
         try:
@@ -60,8 +60,8 @@ def encrypt(self, key, iv="", cek="", **kwargs):
 
         _enc = self["enc"]
         _auth_data = jwe.b64_encode_header()
-        ctxt, tag, cek = self.enc_setup(_enc, _msg.encode(),
-                                        auth_data=_auth_data, key=cek, iv=iv)
+        ctxt, tag, cek = self.enc_setup(_enc, _msg, auth_data=_auth_data,
+                                        key=cek, iv=iv)
         return jwe.pack(parts=[jek, iv, ctxt, tag])
 
     def decrypt(self, token, key=None, cek=None):

tests/test_07_jwe.py

@@ -393,4 +393,15 @@ def test_encrypt_jwk_key():
     assert _enc
     decryptor = factory(_enc, alg="ECDH-ES", enc="A128GCM")
     res = decryptor.decrypt()
-    assert res == plain
+    assert res == plain
+
+
+def test_sym_encrypt_decrypt_JWE():
+    encryption_key = SYMKey(use="enc", key='DukeofHazardpass',
+                            kid="some-key-id")
+    jwe = JWE(plain, alg="A128KW", enc="A128CBC-HS256")
+    _jwe = jwe.encrypt(keys=[encryption_key], kid="some-key-id")
+    decryptor = factory(_jwe, alg="A128KW", enc="A128CBC-HS256")
+
+    resp = decryptor.decrypt(_jwe, [encryption_key])
+    assert resp == plain