Fix array payload issue

Author: e107109

client_encryption/field_level_encryption.py

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

client_encryption/json_path_utils.py

@@ -42,7 +42,12 @@ def update_node(tree, path, node_str):
         except json.JSONDecodeError:
             node_json = node_str
 
-        if to_set in current_node and type(current_node[to_set]) is dict and type(node_json) is dict:
+        if type(current_node) is list:
+            if to_set in current_node[0] and type(current_node[0][to_set]) is dict and type(node_json) is dict:
+                current_node[0][to_set].update(node_json)
+            else:
+                current_node[0][to_set] = node_json
+        elif to_set in current_node and type(current_node[to_set]) is dict and type(node_json) is dict:
             current_node[to_set].update(node_json)
         else:
             current_node[to_set] = node_json
@@ -66,7 +71,10 @@ def pop_node(tree, path):
         else:
             node = tree
 
-        deleted_elem = node.pop(to_delete)
+        if type(node) is list:
+            deleted_elem = node[0].pop(to_delete)
+        else:
+            deleted_elem = node.pop(to_delete)
         if isinstance(deleted_elem, str):
             return deleted_elem
         else:
@@ -91,8 +99,9 @@ def cleanup_node(tree, path, target):
                 node = __get_node(tree, parent, False)
             else:
                 node = tree
-
-            if not node[to_delete]:
+            if type(node) is list and not node[0][to_delete]:
+                del node[0][to_delete]
+            elif not node[to_delete]:
                 del node[to_delete]
 
     else:
@@ -107,12 +116,23 @@ def __get_node(tree, node_list, create):
     last_node = node_list.pop()
 
     for node in node_list:
-        current = current[node]
+        if type(current) is list:
+            current = current[0][node]
+        else:
+            current = current[node]
 
-    if type(current) is not dict:
+    if type(current) is not dict and type(current) is not list:
         raise ValueError("'" + current + "' is not of dict type")
 
-    if last_node not in current and create:
+    if type(current) is list:
+        if not current and create:
+            d = dict()
+            d[last_node] = {}
+            current.append(d)
+        elif last_node not in current[0] and create:
+            current[0][last_node] = {}
+        return current[0][last_node]
+    elif last_node not in current and create:
         current[last_node] = {}
 
     return current[last_node]

tests/test_field_level_encryption.py

@@ -78,6 +78,20 @@ def __assert_payload_encrypted(self, payload, encrypted, config):
         del payload["encryptedData"]
         self.assertEqual(payload, to_test.decrypt_payload(encrypted, config))
 
+    def __assert_array_payload_encrypted(self, payload, encrypted, config):
+        self.assertNotIn("data", encrypted[0])
+        self.assertIn("encryptedData", encrypted[0])
+        enc_data = encrypted[0]["encryptedData"]
+        self.assertEqual(6, len(enc_data.keys()))
+        self.assertIsNotNone(enc_data["iv"])
+        self.assertIsNotNone(enc_data["encryptedKey"])
+        self.assertIsNotNone(enc_data["encryptedValue"])
+        self.assertEqual("SHA256", enc_data["oaepHashingAlgo"])
+        self.assertEqual("761b003c1eade3a5490e5000d37887baa5e6ec0e226c07706e599451fc032a79", enc_data["keyFingerprint"])
+        self.assertEqual("80810fc13a8319fcf0e2ec322c82a4c304b782cc3ce671176343cfe8160c2279", enc_data["certFingerprint"])
+        del payload[0]["encryptedData"]
+        self.assertEqual(payload, to_test.decrypt_payload(encrypted, config))
+
     def test_encrypt_payload_base64_field_encoding(self):
         payload = {
             "data": {
@@ -155,6 +169,155 @@ def test_encrypt_payload_with_type_list(self):
         encrypted_payload = to_test.encrypt_payload(payload, self._config)
         self.__assert_payload_encrypted(payload, encrypted_payload, self._config)
 
+    def test_encrypt_array_payload_with_type_string(self):
+        payload = [{
+            "data": "item1",
+            "encryptedData": {}
+        }]
+
+        encrypted_payload = to_test.encrypt_payload(payload, self._config)
+        self.__assert_array_payload_encrypted(payload, encrypted_payload, self._config)
+
+    def test_encrypt_array_payload_with_type_list(self):
+        payload = [{
+            "data": ["item1", "item2", "item3"],
+            "encryptedData": {}
+        }]
+
+        encrypted_payload = to_test.encrypt_payload(payload, self._config)
+        self.__assert_array_payload_encrypted(payload, encrypted_payload, self._config)
+
+    def test_encrypt_array_payload_with_type_object(self):
+
+        payload = [{
+            "data": {
+                "field1": "value1",
+                "field2": "value2"
+            },
+            "encryptedData": {}
+        }]
+
+        encrypted_payload = to_test.encrypt_payload(payload, self._config)
+        self.__assert_array_payload_encrypted(payload, encrypted_payload, self._config)
+
+    def test_encrypt_array_payload_with_type_multiple_object(self):
+
+        payload = [{
+            "data": {
+                "field1": "value1",
+                "field2": "value2"
+            },
+            "encryptedData": {}
+            },
+            {
+                "data": {
+                    "field1": "value1",
+                    "field2": "value2"
+                },
+                "encryptedData": {}
+            }
+        ]
+
+        encrypted_payload = to_test.encrypt_payload(payload, self._config)
+        self.__assert_array_payload_encrypted(payload, encrypted_payload, self._config)
+
+    def test_encrypt_array_payload_skip_when_in_path_does_not_exist(self):
+        payload = [{
+            "dataNotToEncrypt": {
+                "field1": "value1",
+                "field2": "value2"
+            },
+            "encryptedData": {}
+        }]
+
+        encrypted_payload = to_test.encrypt_payload(payload, self._config)
+
+        self.assertEqual(payload, encrypted_payload)
+
+    def test_encrypt_array_payload_create_node_when_out_path_parent_exists(self):
+        self._config._paths["$"]._to_encrypt = {"data": "encryptedDataParent.encryptedData"}
+
+        payload = [{
+            "data": {
+                "field1": "value1",
+                "field2": "value2"
+            },
+            "encryptedDataParent": {}
+        }]
+
+        encrypted_payload = to_test.encrypt_payload(payload, self._config)
+
+        self.assertNotIn("data", encrypted_payload[0])
+        self.assertIn("encryptedDataParent", encrypted_payload[0])
+        self.assertIn("encryptedData", encrypted_payload[0]["encryptedDataParent"])
+
+    def test_encrypt_array_payload_with_multiple_encryption_paths(self):
+        self._config._paths["$"]._to_encrypt = {"data1": "encryptedData1", "data2": "encryptedData2"}
+
+        payload = [{
+            "data1": {
+                "field1": "value1",
+                "field2": "value2"
+            },
+            "data2": {
+                "field3": "value3",
+                "field4": "value4"
+            },
+            "encryptedData1": {},
+            "encryptedData2": {}
+        }]
+
+        encrypted_payload = to_test.encrypt_payload(payload, self._config)
+
+        self.assertNotIn("data1", encrypted_payload[0])
+        self.assertNotIn("data2", encrypted_payload[0])
+        enc_data1 = encrypted_payload[0]["encryptedData1"]
+        enc_data2 = encrypted_payload[0]["encryptedData2"]
+        self.assertIsNotNone(enc_data1["iv"])
+        self.assertIsNotNone(enc_data1["encryptedKey"])
+        self.assertIsNotNone(enc_data1["encryptedValue"])
+        self.assertIsNotNone(enc_data2["iv"])
+        self.assertIsNotNone(enc_data2["encryptedKey"])
+        self.assertIsNotNone(enc_data2["encryptedValue"])
+        self.assertNotEqual(enc_data1["iv"], enc_data2["iv"], "using same set of params")
+
+    def test_encrypt_array_payload_when_root_as_in_path(self):
+        self._config._paths["$"]._to_encrypt = {"$": "encryptedData"}
+
+        payload = [{
+            "field1": "value1",
+            "field2": "value2"
+        }]
+
+        encrypted_payload = to_test.encrypt_payload(payload, self._config)
+
+        self.assertNotIn("field1", encrypted_payload[0])
+        self.assertNotIn("field2", encrypted_payload[0])
+        self.assertIn("encryptedData", encrypted_payload[0])
+        self.assertEqual(6, len(encrypted_payload[0]["encryptedData"].keys()))
+
+    def test_encrypt_array_payload_when_out_path_same_as_in_path(self):
+        self._config._paths["$"]._to_encrypt = {"data": "data"}
+
+        payload = [{
+            "data": {
+                "field1": "value1",
+                "field2": "value2"
+            }
+        }]
+
+        encrypted_payload = to_test.encrypt_payload(payload, self._config)
+
+        self.assertIn("data", encrypted_payload[0])
+        self.assertNotIn("field1", encrypted_payload[0]["data"])
+        self.assertNotIn("field2", encrypted_payload[0]["data"])
+        self.assertIn("iv", encrypted_payload[0]["data"])
+        self.assertIn("encryptedKey", encrypted_payload[0]["data"])
+        self.assertIn("encryptedValue", encrypted_payload[0]["data"])
+        self.assertIn("certFingerprint", encrypted_payload[0]["data"])
+        self.assertIn("keyFingerprint", encrypted_payload[0]["data"])
+        self.assertIn("oaepHashingAlgo", encrypted_payload[0]["data"])
+
     def test_encrypt_payload_skip_when_in_path_does_not_exist(self):
         payload = {
             "dataNotToEncrypt": {