Merge branch 'master' of https://github.com/rohe/pyjwkest

Author: rohe

requirements.txt

@@ -1,2 +1,2 @@
 requests
-pycrypto >= 2.2
+pycryptodomex>=3.4.2

src/jwkest/jwe.py

@@ -489,15 +489,20 @@ def encrypt(self, key, iv="", cek="", **kwargs):
             else:
                 raise ParameterError("Zip has unknown value: %s" % self["zip"])
 
+        kwarg_cek = cek or None
+
         _enc = self["enc"]
         cek, iv = self._generate_key_and_iv(_enc, cek, iv)
+        self["cek"] = cek
 
         logger.debug("cek: %s, iv: %s" % ([c for c in cek], [c for c in iv]))
 
         _encrypt = RSAEncrypter(self.with_digest).encrypt
 
         _alg = self["alg"]
-        if _alg == "RSA-OAEP":
+        if kwarg_cek:
+            jwe_enc_key = ''
+        elif _alg == "RSA-OAEP":
             jwe_enc_key = _encrypt(cek, key, 'pkcs1_oaep_padding')
         elif _alg == "RSA1_5":
             jwe_enc_key = _encrypt(cek, key)
@@ -511,7 +516,7 @@ def encrypt(self, key, iv="", cek="", **kwargs):
         ctxt, tag, key = self.enc_setup(_enc, _msg, enc_header, cek, iv)
         return jwe.pack(parts=[jwe_enc_key, iv, ctxt, tag])
 
-    def decrypt(self, token, key):
+    def decrypt(self, token, key, cek=None):
         """ Decrypts a JWT
 
         :param token: The JWT
@@ -529,13 +534,16 @@ def decrypt(self, token, key):
         _decrypt = RSAEncrypter(self.with_digest).decrypt
 
         _alg = jwe.headers["alg"]
-        if _alg == "RSA-OAEP":
+        if cek:
+            pass
+        elif _alg == "RSA-OAEP":
             cek = _decrypt(jek, key, 'pkcs1_oaep_padding')
         elif _alg == "RSA1_5":
             cek = _decrypt(jek, key)
         else:
             raise NotSupportedAlgorithm(_alg)
 
+        self["cek"] = cek
         enc = jwe.headers["enc"]
         try:
             assert enc in SUPPORTED["enc"]
@@ -687,7 +695,7 @@ def encrypt(self, key, iv="", cek="", **kwargs):
             return jwe.pack(parts=[kwargs['encrypted_key'], iv, ctxt, tag])
         return jwe.pack(parts=[iv, ctxt, tag])
 
-    def decrypt(self, token=None, key=None):
+    def decrypt(self, token=None, key=None, **kwargs):
 
         if not self.cek:
             raise Exception("Content Encryption Key is Not Yet Set")
@@ -747,7 +755,7 @@ def encrypt(self, keys=None, cek="", iv="", **kwargs):
         :return: Encrypted message
         """
 
-        encrypted_key = cek = iv = None
+        # encrypted_key = cek = iv = None
         _alg = self["alg"]
 
         # Find Usable Keys
@@ -801,6 +809,7 @@ def encrypt(self, keys=None, cek="", iv="", **kwargs):
 
             try:
                 token = encrypter.encrypt(_key, **kwargs)
+                self["cek"] = encrypter.cek if 'cek' in encrypter else None
             except TypeError as err:
                 raise err
             else:
@@ -811,7 +820,7 @@ def encrypt(self, keys=None, cek="", iv="", **kwargs):
         logger.error("Could not find any suitable encryption key")
         raise NoSuitableEncryptionKey()
 
-    def decrypt(self, token=None, keys=None, alg=None):
+    def decrypt(self, token=None, keys=None, alg=None, cek=None):
         if token:
             jwe = JWEnc().unpack(token)
             # header, ek, eiv, ctxt, tag = token.split(b".")
@@ -829,7 +838,7 @@ def decrypt(self, token=None, keys=None, alg=None):
         else:
             keys = self._pick_keys(self._get_keys(), use="enc", alg=_alg)
 
-        if not keys:
+        if not keys and not cek:
             raise NoSuitableDecryptionKey(_alg)
 
         if _alg in ["RSA-OAEP", "RSA1_5"]:
@@ -847,10 +856,21 @@ def decrypt(self, token=None, keys=None, alg=None):
         else:
             raise NotSupportedAlgorithm
 
+        if cek:
+            try:
+                msg = decrypter.decrypt(as_bytes(token), None, cek=cek)
+                self["cek"] = decrypter.cek if 'cek' in decrypter else None
+            except (KeyError, DecryptionFailed):
+                pass
+            else:
+                logger.debug("Decrypted message using exiting CEK")
+                return msg
+
         for key in keys:
             _key = key.encryption_key(alg=_alg, private=False)
             try:
                 msg = decrypter.decrypt(as_bytes(token), _key)
+                self["cek"] = decrypter.cek if 'cek' in decrypter else None
             except (KeyError, DecryptionFailed):
                 pass
             else:

src/jwkest/jwk.py

@@ -22,6 +22,7 @@
 from jwkest import JWKESTException
 from jwkest import b64d
 from jwkest import b64e
+from jwkest import UnknownAlgorithm
 from jwkest.ecc import NISTEllipticCurve
 from jwkest.jwt import b2s_conv
 
@@ -87,18 +88,18 @@ def sha512_digest(msg):
 # =============================================================================
 
 
-def import_rsa_key_from_file(filename):
-    return RSA.importKey(open(filename, 'r').read())
+def import_rsa_key_from_file(filename, passphrase=None):
+    return RSA.importKey(open(filename, 'r').read(), passphrase=passphrase)
 
 
-def import_rsa_key(key):
+def import_rsa_key(key, passphrase=None):
     """
     Extract an RSA key from a PEM-encoded certificate
-
     :param key: RSA key encoded in standard form
+    :param passphrase: Password to open the certificate (Optional)
     :return: RSA key instance
     """
-    return importKey(key)
+    return importKey(key, passphrase=passphrase)
 
 
 def der2rsa(der):
@@ -189,6 +190,39 @@ def x509_rsa_load(txt):
     return [("rsa", import_rsa_key(txt))]
 
 
+def key_from_jwk_dict(jwk_dict, private=True):
+    """Load JWK from dictionary"""
+    if jwk_dict['kty'] == 'EC':
+        if private:
+            return ECKey(kid=jwk_dict['kid'],
+                         crv=jwk_dict['crv'],
+                         x=jwk_dict['x'],
+                         y=jwk_dict['y'],
+                         d=jwk_dict['d'])
+        else:
+            return ECKey(kid=jwk_dict['kid'],
+                         crv=jwk_dict['crv'],
+                         x=jwk_dict['x'],
+                         y=jwk_dict['y'])
+    elif jwk_dict['kty'] == 'RSA':
+        if private:
+            return RSAKey(kid=jwk_dict['kid'],
+                          n=jwk_dict['n'],
+                          e=jwk_dict['e'],
+                          d=jwk_dict['d'],
+                          p=jwk_dict['p'],
+                          q=jwk_dict['q'])
+        else:
+            return RSAKey(kid=jwk_dict['kid'],
+                          n=jwk_dict['n'],
+                          e=jwk_dict['e'])
+    elif jwk_dict['kty'] == 'oct':
+        return SYMKey(kid=jwk_dict['kid'],
+                      k=jwk_dict['k'])
+    else:
+        raise UnknownAlgorithm
+
+
 class Key(object):
     """
     Basic JSON Web key class

src/jwkest/jws.py

@@ -161,7 +161,10 @@ def sign(self, msg, key):
 
     def verify(self, msg, sig, key):
         h = bytes_to_long(self.digest.new(msg).digest())
-        return self._sign.verify(h, sig, key)
+        if self._sign.verify(h, sig, key):
+            return True
+        else:
+            raise BadSignature()
 
 
 class PSSSigner(Signer):

tests/test_4_jwe.py

@@ -193,6 +193,37 @@ def full_path(local_file):
 rsa = RSA.importKey(open(KEY, 'r').read())
 plain = b'Now is the time for all good men to come to the aid of their country.'
 
+def test_cek_reuse_encryption_rsaes_rsa15():
+
+    _rsa = JWE_RSA(plain, alg="RSA1_5", enc="A128CBC-HS256")
+    jwt = _rsa.encrypt(rsa)
+    dec = JWE_RSA()
+    msg = dec.decrypt(jwt, rsa)
+
+    assert msg == plain
+
+    _rsa2 = JWE_RSA(plain, alg="RSA1_5", enc="A128CBC-HS256")
+    jwt = _rsa2.encrypt(None, cek=dec["cek"])
+    dec2 = JWE_RSA()
+    msg = dec2.decrypt(jwt, None, cek=_rsa["cek"])
+
+    assert msg == plain
+
+def test_cek_reuse_encryption_rsaes_rsa_oaep():
+
+    _rsa = JWE_RSA(plain, alg="RSA-OAEP", enc="A256GCM")
+    jwt = _rsa.encrypt(rsa)
+    dec = JWE_RSA()
+    msg = dec.decrypt(jwt, rsa)
+
+    assert msg == plain
+
+    _rsa2 = JWE_RSA(plain, alg="RSA-OAEP", enc="A256GCM")
+    jwt = _rsa2.encrypt(None, cek=dec["cek"])
+    dec2 = JWE_RSA()
+    msg = dec2.decrypt(jwt, None, cek=_rsa["cek"])
+
+    assert msg == plain
 
 def test_rsa_encrypt_decrypt_rsa_cbc():
     _rsa = JWE_RSA(plain, alg="RSA1_5", enc="A128CBC-HS256")