fixed oaep algorithm length issue

moved decorator to 'call_api' function

Author: pqlMC

client_encryption/api_encryption.py

@@ -18,25 +18,24 @@ def __init__(self, encryption_conf_file):
                 self._encryption_conf = FieldLevelEncryptionConfig(json_file.read())
 
     def field_encryption(self, func):
-        """Decorator for API request. func is APIClient.request"""
+        """Decorator for API call_api. func is APIClient.call_api"""
 
         @wraps(func)
-        def request_function(*args, **kwargs):
-            """Wrap request and add field encryption layer to it."""
+        def call_api_function(*args, **kwargs):
+            """Wrap call_api and add field encryption layer to it."""
 
             in_body = kwargs.get("body", None)
             kwargs["body"] = self._encrypt_payload(kwargs.get("headers", None), in_body) if in_body else in_body
+            kwargs["_preload_content"] = False
 
             response = func(*args, **kwargs)
 
-            if type(response.data) is not str:
-                response_body = self._decrypt_payload(response.getheaders(), response.json())
-                response._content = json.dumps(response_body, indent=4).encode('utf-8')
+            response_body = self._decrypt_payload(response.getheaders(), response.data)
 
-            return response
+            return response_body
 
-        request_function.__fle__ = True
-        return request_function
+        call_api_function.__fle__ = True
+        return call_api_function
 
     def _encrypt_payload(self, headers, body):
         """Encryption enforcement based on configuration - encrypt and add session key params to header or body"""
@@ -86,10 +85,10 @@ def _decrypt_payload(self, headers, body):
 
 
 def add_encryption_layer(api_client, encryption_conf_file):
-    """Decorate APIClient.request with field level encryption"""
+    """Decorate APIClient.call_api with field level encryption"""
 
     api_encryption = ApiEncryption(encryption_conf_file)
-    api_client.request = api_encryption.field_encryption(api_client.request)
+    api_client.call_api = api_encryption.field_encryption(api_client.call_api)
 
     __check_oauth(api_client)  # warn the user if authentication layer is missing/not set
 

client_encryption/encryption_utils.py

@@ -54,15 +54,21 @@ def __get_crypto_file_type(file_content):
         return FILETYPE_ASN1
 
 
-def load_hash_algorithm(algo_str):
-    """Load a hash algorithm object of Crypto.Hash from a list of supported ones."""
+def validate_hash_algorithm(algo_str):
+    """Validate a hash algorithm against a list of supported ones."""
 
     if algo_str:
         algo_key = algo_str.replace("-", "").upper()
 
         if algo_key in _SUPPORTED_HASH:
-            return _SUPPORTED_HASH[algo_key]
+            return algo_key
         else:
             raise HashAlgorithmError("Hash algorithm invalid or not supported.")
     else:
         raise HashAlgorithmError("No hash algorithm provided.")
+
+
+def load_hash_algorithm(algo_str):
+    """Load a hash algorithm object of Crypto.Hash from a list of supported ones."""
+
+    return _SUPPORTED_HASH[validate_hash_algorithm(algo_str)]

client_encryption/field_level_encryption.py

@@ -12,10 +12,10 @@ def encrypt_payload(payload, config, _params=None):
     """Encrypt some fields of a JSON payload using the given configuration."""
 
     try:
-        if type(payload) is str:
-            json_payload = json.loads(payload)
-        else:
+        if type(payload) is dict:
             json_payload = copy.deepcopy(payload)
+        else:
+            json_payload = json.loads(payload)
 
         for elem, target in config.paths["$"].to_encrypt.items():
             if not _params:
@@ -50,10 +50,10 @@ def decrypt_payload(payload, config, _params=None):
     """Decrypt some fields of a JSON payload using the given configuration."""
 
     try:
-        if type(payload) is str:
-            json_payload = json.loads(payload)
-        else:
+        if type(payload) is dict:
             json_payload = payload
+        else:
+            json_payload = json.loads(payload)
 
         for elem, target in config.paths["$"].to_decrypt.items():
             try:

client_encryption/field_level_encryption_config.py

@@ -2,7 +2,7 @@
 from OpenSSL.crypto import dump_certificate, FILETYPE_ASN1, dump_publickey
 from Crypto.Hash import SHA256
 from client_encryption.encoding_utils import Encoding
-from client_encryption.encryption_utils import load_encryption_certificate, load_decryption_key, load_hash_algorithm
+from client_encryption.encryption_utils import load_encryption_certificate, load_decryption_key, validate_hash_algorithm
 
 
 class FieldLevelEncryptionConfig(object):
@@ -44,9 +44,7 @@ def __init__(self, conf):
         else:
             self._decryption_key = None
 
-        digest_algo = json_config["oaepPaddingDigestAlgorithm"]
-        if load_hash_algorithm(digest_algo) is not None:
-            self._oaep_padding_digest_algorithm = digest_algo
+        self._oaep_padding_digest_algorithm = validate_hash_algorithm(json_config["oaepPaddingDigestAlgorithm"])
 
         data_enc = Encoding(json_config["dataEncoding"].upper())
         self._data_encoding = data_enc

client_encryption/session_key_params.py

@@ -1,6 +1,6 @@
 from binascii import Error
 from Crypto.Cipher import PKCS1_OAEP, AES
-from Crypto import Random
+from Crypto.Random import get_random_bytes
 from Crypto.PublicKey import RSA
 from client_encryption.encoding_utils import encode_bytes, decode_value
 from client_encryption.encryption_utils import load_hash_algorithm
@@ -62,11 +62,11 @@ def generate(config):
         encoding = config.data_encoding
 
         # Generate a random IV
-        iv = Random.new().read(SessionKeyParams._BLOCK_SIZE)
+        iv = get_random_bytes(SessionKeyParams._BLOCK_SIZE)
         iv_encoded = encode_bytes(iv, encoding)
 
         # Generate an AES secret key
-        secret_key = Random.new().read(SessionKeyParams._KEY_SIZE)
+        secret_key = get_random_bytes(SessionKeyParams._KEY_SIZE)
 
         # Encrypt the secret key
         secret_key_encrypted = SessionKeyParams.__wrap_secret_key(secret_key, config)

client_encryption/version.py

@@ -1,3 +1,3 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
-__version__ = "1.0.2"
+__version__ = "1.0.3-dev"

tests/test_encryption_utils.py

@@ -131,6 +131,30 @@ def test_load_hash_algorithm_underscore(self):
     def test_load_hash_algorithm_none(self):
         self.assertRaises(HashAlgorithmError, to_test.load_hash_algorithm, None)
 
+    def test_validate_hash_algorithm(self):
+        hash_algo = to_test.validate_hash_algorithm("SHA224")
+
+        self.assertEqual(hash_algo, "SHA224")
+
+    def test_validate_hash_algorithm_dash(self):
+        hash_algo = to_test.validate_hash_algorithm("SHA-512")
+
+        self.assertEqual(hash_algo, "SHA512")
+
+    def test_validate_hash_algorithm_lowercase(self):
+        hash_algo = to_test.validate_hash_algorithm("sha384")
+
+        self.assertEqual(hash_algo, "SHA384")
+
+    def test_validate_hash_algorithm_not_supported(self):
+        self.assertRaises(HashAlgorithmError, to_test.validate_hash_algorithm, "MD5")
+
+    def test_validate_hash_algorithm_underscore(self):
+        self.assertRaises(HashAlgorithmError, to_test.validate_hash_algorithm, "SHA_512")
+
+    def test_validate_hash_algorithm_none(self):
+        self.assertRaises(HashAlgorithmError, to_test.validate_hash_algorithm, None)
+
     @staticmethod
     def __strip_key(rsa_key):
         return rsa_key.export_key(pkcs=8).decode('utf-8').replace("\n", "")[27:-25]

tests/test_field_level_encryption_config.py

@@ -151,7 +151,7 @@ def test_load_config_SHA512_oaep_padding_algorithm(self):
         json_conf["oaepPaddingDigestAlgorithm"] = oaep_algo_test
 
         conf = to_test.FieldLevelEncryptionConfig(json_conf)
-        self.__check_configuration(conf, oaep_algo=oaep_algo_test)
+        self.__check_configuration(conf, oaep_algo="SHA512")
 
     def test_load_config_wrong_oaep_padding_algorithm(self):
         oaep_algo_test = "sha_512"