Add AESKey and DIRKey implementations

Author: blag

jose/backends/__init__.py

@@ -29,7 +29,17 @@
 except ImportError:
     from jose.backends.ecdsa_backend import ECDSAECKey as ECKey  # noqa: F401
 
+try:
+    from jose.backends.cryptography_backend import CryptographyAESKey as AESKey  # noqa: F401
+except ImportError:
+    try:
+        from jose.backends.pycrypto_backend import AESKey  # noqa: F401
+    except ImportError:
+        AESKey = None
+
 try:
     from jose.backends.cryptography_backend import CryptographyHMACKey as HMACKey  # noqa: F401
 except ImportError:
     from jose.backends.native import HMACKey  # noqa: F401
+
+from .base import DIRKey  # noqa: F401

jose/backends/base.py

@@ -1,3 +1,8 @@
+import six
+
+from ..utils import base64url_encode
+
+
 class Key(object):
     """
     A simple interface for implementing JWK keys.
@@ -19,3 +24,67 @@ def to_pem(self):
 
     def to_dict(self):
         raise NotImplementedError()
+
+    def encrypt(self, plain_text, aad=None):
+        """
+        Encrypt the plain text and generate an auth tag if appropriate
+
+        Args:
+            plain_text (bytes): Data to encrypt
+            aad (bytes, optional): Authenticated Additional Data if key's algorithm supports auth mode
+
+        Returns:
+            (bytes, bytes, bytes): IV, cipher text, and auth tag
+        """
+        raise NotImplementedError()
+
+    def decrypt(self, cipher_text, iv=None, aad=None, tag=None):
+        """
+        Decrypt the cipher text and validate the auth tag if present
+        Args:
+            cipher_text (bytes): Cipher text to decrypt
+            iv (bytes): IV if block mode
+            aad (bytes): Additional Authenticated Data to verify if auth mode
+            tag (bytes): Authentication tag if auth mode
+
+        Returns:
+            bytes: Decrypted value
+        """
+        raise NotImplementedError()
+
+    def wrap_key(self, key_data):
+        """
+        Wrap the the plain text key data
+
+        Args:
+            key_data (bytes): Key data to wrap
+
+        Returns:
+            bytes: Wrapped key
+        """
+        raise NotImplementedError()
+
+    def unwrap_key(self, wrapped_key):
+        """
+        Unwrap the the wrapped key data
+
+        Args:
+            wrapped_key (bytes): Wrapped key data to unwrap
+
+        Returns:
+            bytes: Unwrapped key
+        """
+        raise NotImplementedError()
+
+
+class DIRKey(Key):
+    def __init__(self, key_data, algorithm):
+        self._key = six.ensure_binary(key_data)
+        self._alg = algorithm
+
+    def to_dict(self):
+        return {
+            'alg': self._alg,
+            'kty': 'oct',
+            'k': base64url_encode(self._key),
+        }

jose/backends/cryptography_backend.py

@@ -8,14 +8,17 @@
 from .base import Key
 from ..utils import base64_to_long, long_to_base64, base64url_decode, base64url_encode
 from ..constants import ALGORITHMS
-from ..exceptions import JWKError
+from ..exceptions import JWKError, JWEError
 
-from cryptography.exceptions import InvalidSignature
+from cryptography.exceptions import InvalidSignature, InvalidTag
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.bindings.openssl.binding import Binding
 from cryptography.hazmat.primitives import hashes, serialization, hmac
 from cryptography.hazmat.primitives.asymmetric import ec, rsa, padding
 from cryptography.hazmat.primitives.asymmetric.utils import decode_dss_signature, encode_dss_signature
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, aead, modes
+from cryptography.hazmat.primitives.keywrap import aes_key_wrap, aes_key_unwrap, InvalidUnwrap
+from cryptography.hazmat.primitives.padding import PKCS7
 from cryptography.hazmat.primitives.serialization import load_pem_private_key, load_pem_public_key
 from cryptography.utils import int_from_bytes, int_to_bytes
 from cryptography.x509 import load_pem_x509_certificate
@@ -223,6 +226,10 @@ class CryptographyRSAKey(Key):
     SHA384 = hashes.SHA384
     SHA512 = hashes.SHA512
 
+    RSA1_5 = padding.PKCS1v15()
+    RSA_OAEP = padding.OAEP(padding.MGF1(hashes.SHA1()), hashes.SHA1(), None)
+    RSA_OAEP_256 = padding.OAEP(padding.MGF1(hashes.SHA256()), hashes.SHA256(), None)
+
     def __init__(self, key, algorithm, cryptography_backend=default_backend):
         if algorithm not in ALGORITHMS.RSA:
             raise JWKError('hash_alg: %s is not a valid hash algorithm' % algorithm)
@@ -234,6 +241,12 @@ def __init__(self, key, algorithm, cryptography_backend=default_backend):
         }.get(algorithm)
         self._algorithm = algorithm
 
+        self.padding = {
+            ALGORITHMS.RSA1_5: self.RSA1_5,
+            ALGORITHMS.RSA_OAEP: self.RSA_OAEP,
+            ALGORITHMS.RSA_OAEP_256: self.RSA_OAEP_256
+        }.get(algorithm)
+
         self.cryptography_backend = cryptography_backend
 
         # if it conforms to RSAPublicKey interface
@@ -396,6 +409,149 @@ def to_dict(self):
 
         return data
 
+    def wrap_key(self, key_data):
+        try:
+            wrapped_key = self.prepared_key.encrypt(key_data, self.padding)
+        except Exception as e:
+            raise JWEError(e)
+
+        return wrapped_key
+
+    def unwrap_key(self, wrapped_key):
+        try:
+            unwrapped_key = self.prepared_key.decrypt(
+                wrapped_key,
+                self.padding
+            )
+            return unwrapped_key
+        except Exception as e:
+            raise JWEError(e)
+
+
+class CryptographyAESKey(Key):
+    KEY_128 = (ALGORITHMS.A128GCM, ALGORITHMS.A128GCMKW, ALGORITHMS.A128KW,
+               ALGORITHMS.A128CBC)
+    KEY_192 = (ALGORITHMS.A192GCM, ALGORITHMS.A192GCMKW, ALGORITHMS.A192KW,
+               ALGORITHMS.A192CBC)
+    KEY_256 = (ALGORITHMS.A256GCM, ALGORITHMS.A256GCMKW, ALGORITHMS.A256KW,
+               ALGORITHMS.A128CBC_HS256, ALGORITHMS.A256CBC)
+    KEY_384 = (ALGORITHMS.A192CBC_HS384,)
+    KEY_512 = (ALGORITHMS.A256CBC_HS512,)
+
+    AES_KW_ALGS = (ALGORITHMS.A128KW, ALGORITHMS.A192KW, ALGORITHMS.A256KW)
+
+    MODES = {
+        ALGORITHMS.A128GCM: modes.GCM,
+        ALGORITHMS.A192GCM: modes.GCM,
+        ALGORITHMS.A256GCM: modes.GCM,
+        ALGORITHMS.A128CBC_HS256: modes.CBC,
+        ALGORITHMS.A192CBC_HS384: modes.CBC,
+        ALGORITHMS.A256CBC_HS512: modes.CBC,
+        ALGORITHMS.A128CBC: modes.CBC,
+        ALGORITHMS.A192CBC: modes.CBC,
+        ALGORITHMS.A256CBC: modes.CBC,
+        ALGORITHMS.A128GCMKW: modes.GCM,
+        ALGORITHMS.A192GCMKW: modes.GCM,
+        ALGORITHMS.A256GCMKW: modes.GCM,
+        ALGORITHMS.A128KW: None,
+        ALGORITHMS.A192KW: None,
+        ALGORITHMS.A256KW: None
+    }
+
+    def __init__(self, key, algorithm):
+        if algorithm not in ALGORITHMS.AES:
+            raise JWKError('%s is not a valid AES algorithm' % algorithm)
+        if algorithm not in ALGORITHMS.SUPPORTED.union(ALGORITHMS.AES_PSEUDO):
+            raise JWKError('%s is not a supported algorithm' % algorithm)
+
+        self._algorithm = algorithm
+        self._mode = self.MODES.get(self._algorithm)
+
+        if algorithm in self.KEY_128 and len(key) != 16:
+            raise JWKError("Key must be 128 bit for alg {}".format(algorithm))
+        elif algorithm in self.KEY_192 and len(key) != 24:
+            raise JWKError("Key must be 192 bit for alg {}".format(algorithm))
+        elif algorithm in self.KEY_256 and len(key) != 32:
+            raise JWKError("Key must be 256 bit for alg {}".format(algorithm))
+        elif algorithm in self.KEY_384 and len(key) != 48:
+            raise JWKError("Key must be 384 bit for alg {}".format(algorithm))
+        elif algorithm in self.KEY_512 and len(key) != 64:
+            raise JWKError("Key must be 512 bit for alg {}".format(algorithm))
+
+        self._key = key
+
+    def to_dict(self):
+        data = {
+            'alg': self._algorithm,
+            'kty': 'oct',
+            'k': base64url_encode(self._key)
+        }
+        return data
+
+    def encrypt(self, plain_text, aad=None):
+        plain_text = six.ensure_binary(plain_text)
+        try:
+            iv = get_random_bytes(algorithms.AES.block_size//8)
+            mode = self._mode(iv)
+            if mode.name == "GCM":
+                cipher = aead.AESGCM(self._key)
+                cipher_text_and_tag = cipher.encrypt(iv, plain_text, aad)
+                cipher_text = cipher_text_and_tag[:len(cipher_text_and_tag) - 16]
+                auth_tag = cipher_text_and_tag[-16:]
+            else:
+                cipher = Cipher(algorithms.AES(self._key), mode,
+                                backend=default_backend())
+                encryptor = cipher.encryptor()
+                padder = PKCS7(algorithms.AES.block_size).padder()
+                padded_data = padder.update(plain_text)
+                padded_data += padder.finalize()
+                cipher_text = encryptor.update(padded_data) + encryptor.finalize()
+                auth_tag = None
+            return iv, cipher_text, auth_tag
+        except Exception as e:
+            raise JWEError(e)
+
+    def decrypt(self, cipher_text, iv=None, aad=None, tag=None):
+        cipher_text = six.ensure_binary(cipher_text)
+        try:
+            iv = six.ensure_binary(iv)
+            mode = self._mode(iv)
+            if mode.name == "GCM":
+                if tag is None:
+                    raise ValueError("tag cannot be None")
+                cipher = aead.AESGCM(self._key)
+                cipher_text_and_tag = cipher_text + tag
+                try:
+                    plain_text = cipher.decrypt(iv, cipher_text_and_tag, aad)
+                except InvalidTag:
+                    raise JWEError("Invalid JWE Auth Tag")
+            else:
+                cipher = Cipher(algorithms.AES(self._key), mode,
+                                backend=default_backend())
+                decryptor = cipher.decryptor()
+                padded_plain_text = decryptor.update(cipher_text)
+                padded_plain_text += decryptor.finalize()
+                unpadder = PKCS7(algorithms.AES.block_size).unpadder()
+                plain_text = unpadder.update(padded_plain_text)
+                plain_text += unpadder.finalize()
+
+            return plain_text
+        except Exception as e:
+            raise JWEError(e)
+
+    def wrap_key(self, key_data):
+        key_data = six.ensure_binary(key_data)
+        cipher_text = aes_key_wrap(self._key, key_data, default_backend())
+        return cipher_text  # IV, cipher text, auth tag
+
+    def unwrap_key(self, wrapped_key):
+        wrapped_key = six.ensure_binary(wrapped_key)
+        try:
+            plain_text = aes_key_unwrap(self._key, wrapped_key, default_backend())
+        except InvalidUnwrap as cause:
+            raise JWEError(cause)
+        return plain_text
+
 
 class CryptographyHMACKey(Key):
     """