fix encryption params in header

Author: pqlMC

client_encryption/api_encryption.py

@@ -25,7 +25,7 @@ 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("header_params", None), in_body) if in_body else in_body
+            kwargs["body"] = self._encrypt_payload(args[4], in_body) if in_body else in_body
             kwargs["_preload_content"] = False
 
             response = func(*args, **kwargs)
@@ -44,13 +44,20 @@ def _encrypt_payload(self, headers, body):
         if conf.use_http_headers:
             params = SessionKeyParams.generate(conf)
 
-            headers[conf.iv_field_name] = params.iv_value
-            headers[conf.encrypted_key_field_name] = params.encrypted_key_value
-            headers[conf.encryption_certificate_fingerprint_field_name] = conf.encryption_certificate_fingerprint
-            headers[conf.encryption_key_fingerprint_field_name] = conf.encryption_key_fingerprint
-            headers[conf.oaep_padding_digest_algorithm_field_name] = conf.oaep_padding_digest_algorithm
+            encryption_params = {
+                conf.iv_field_name: params.iv_value,
+                conf.encrypted_key_field_name: params.encrypted_key_value
+            }
+            if conf.encryption_certificate_fingerprint_field_name:
+                encryption_params[conf.encryption_certificate_fingerprint_field_name] = \
+                    conf.encryption_certificate_fingerprint
+            if conf.encryption_key_fingerprint_field_name:
+                encryption_params[conf.encryption_key_fingerprint_field_name] = conf.encryption_key_fingerprint
+            if conf.oaep_padding_digest_algorithm_field_name:
+                encryption_params[conf.oaep_padding_digest_algorithm_field_name] = conf.oaep_padding_digest_algorithm
 
             encrypted_payload = encrypt_payload(body, conf, params)
+            headers.update(encryption_params)
         else:
             encrypted_payload = encrypt_payload(body, conf)
 
@@ -67,10 +74,10 @@ def _decrypt_payload(self, headers, body):
                 iv = headers.pop(conf.iv_field_name)
                 encrypted_key = headers.pop(conf.encrypted_key_field_name)
                 oaep_digest_algo = headers.pop(conf.oaep_padding_digest_algorithm_field_name) \
-                    if conf.oaep_padding_digest_algorithm_field_name in headers else None
-                if conf.encryption_certificate_fingerprint_field_name in headers:
+                    if _contains_param(conf.oaep_padding_digest_algorithm_field_name, headers) else None
+                if _contains_param(conf.encryption_certificate_fingerprint_field_name, headers):
                     del headers[conf.encryption_certificate_fingerprint_field_name]
-                if conf.encryption_key_fingerprint_field_name in headers:
+                if _contains_param(conf.encryption_key_fingerprint_field_name, headers):
                     del headers[conf.encryption_key_fingerprint_field_name]
 
                 params = SessionKeyParams(conf, encrypted_key, iv, oaep_digest_algo)
@@ -84,6 +91,9 @@ def _decrypt_payload(self, headers, body):
         return payload
 
 
+def _contains_param(param_name, headers): return param_name and param_name in headers
+
+
 def add_encryption_layer(api_client, encryption_conf_file):
     """Decorate APIClient.call_api with field level encryption"""
 

tests/utils/api_encryption_test_utils.py

@@ -25,22 +25,22 @@ def __init__(self, api_client=None):
         self.api_client = api_client
 
     def do_something_get(self, **kwargs):
-        return self.api_client.call_api("testservice", "GET", header_params=kwargs["headers"])
+        return self.api_client.call_api("testservice", "GET", None, None, kwargs["headers"])
 
     def do_something_post(self, **kwargs):
-        return self.api_client.call_api("testservice", "POST", header_params=kwargs["headers"], body=kwargs["body"])
+        return self.api_client.call_api("testservice", "POST", None, None, kwargs["headers"], body=kwargs["body"])
 
     def do_something_delete(self, **kwargs):
-        return self.api_client.call_api("testservice", "DELETE", header_params=kwargs["headers"], body=kwargs["body"])
+        return self.api_client.call_api("testservice", "DELETE", None, None, kwargs["headers"], body=kwargs["body"])
 
     def do_something_get_use_headers(self, **kwargs):
-        return self.api_client.call_api("testservice/headers", "GET", header_params=kwargs["headers"])
+        return self.api_client.call_api("testservice/headers", "GET", None, None, kwargs["headers"])
 
     def do_something_post_use_headers(self, **kwargs):
-        return self.api_client.call_api("testservice/headers", "POST", header_params=kwargs["headers"], body=kwargs["body"])
+        return self.api_client.call_api("testservice/headers", "POST", None, None, kwargs["headers"], body=kwargs["body"])
 
     def do_something_delete_use_headers(self, **kwargs):
-        return self.api_client.call_api("testservice/headers", "DELETE", header_params=kwargs["headers"], body=kwargs["body"])
+        return self.api_client.call_api("testservice/headers", "DELETE", None, None, kwargs["headers"], body=kwargs["body"])
 
 
 class MockApiClient(object):