Merge pull request #14 from Mastercard/fix_array_payload

Fix array payload issue

Author: Akshaykdubey

client_encryption/field_level_encryption.py

@@ -12,7 +12,7 @@ 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)
+        json_payload = copy.deepcopy(payload) if type(payload) is dict or type(payload) is list else json.loads(payload)
 
         for elem, target in config.paths["$"].to_encrypt.items():
             if not _params:
@@ -47,7 +47,7 @@ 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)
+        json_payload = payload if type(payload) is dict or type(payload) is list else json.loads(payload)
 
         for elem, target in config.paths["$"].to_decrypt.items():
             try:

client_encryption/json_path_utils.py

@@ -32,17 +32,16 @@ def update_node(tree, path, node_str):
     if __not_root(path):
         parent = path.split(_SEPARATOR)
         to_set = parent.pop()
-        if parent:
-            current_node = __get_node(tree, parent, False)
-        else:
-            current_node = tree
+        current_node = __get_node(tree, parent, False) if parent else tree
 
         try:
             node_json = json.loads(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:
+            update_node_list(to_set, current_node, 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
@@ -53,6 +52,13 @@ def update_node(tree, path, node_str):
     return tree
 
 
+def update_node_list(to_set, current_node, node_json):
+    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
+
+
 def pop_node(tree, path):
     """Retrieve and delete json or value given a path"""
 
@@ -66,7 +72,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 +100,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 +117,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]

requirements.txt

@@ -1,4 +1,5 @@
 pycryptodome==3.8.1
-pyOpenSSL==19.0.0
+pyOpenSSL==22.0.0
 setuptools>=39.0.1
 coverage>=4.5.3
+cryptography==37.0.4

setup.py

@@ -22,5 +22,5 @@
           'Topic :: Software Development :: Libraries :: Python Modules'
       ],
       tests_require=['coverage'],
-      install_requires=['pycryptodome>=3.8.1', 'pyOpenSSL>=19.0.0', 'setuptools>=39.0.1']
+      install_requires=['pycryptodome>=3.8.1', 'pyOpenSSL>=22.0.0', 'setuptools>=39.0.1', 'cryptography==37.0.4']
       )

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": {

tests/test_json_path_utils.py

@@ -21,6 +21,23 @@ def __get_sample_json():
             }
         }
 
+    @staticmethod
+    def __get_array_sample_json():
+        return {
+            "node1": [
+                {
+                    "node2": {
+                    "colour": "red",
+                    "shape": "circle",
+                    "position": {
+                        "lat": 1,
+                        "long": 3
+                        }
+                    }
+                }
+            ]
+        }
+
     def test_get_node(self):
         sample_json = self.__get_sample_json()
 
@@ -151,6 +168,29 @@ def test_update_node_not_json(self):
 
         self.assertIsInstance(node["node1"]["node2"], str, "not a json string")
 
+    def test_update_node_array_with_str(self):
+        sample_json = self.__get_array_sample_json()
+        node = to_test.update_node(sample_json, "node1.node2", "not a json string")
+
+        self.assertIsInstance(node["node1"][0]["node2"], str, "not a json string")
+
+    def test_update_node_array_with_json_str(self):
+        sample_json = self.__get_array_sample_json()
+        node = to_test.update_node(sample_json, "node1.node2", '{"position": {"brightness": 6}}')
+
+        self.assertIsInstance(node["node1"][0]["node2"]["position"], dict)
+        self.assertDictEqual({'node1': [
+                                {'node2': {
+                                    'colour': 'red',
+                                    'shape': 'circle',
+                                    'position': {
+                                        'brightness': 6
+                                    }
+                                }
+                                }
+                            ]}, node)
+
+
     def test_update_node_primitive_type(self):
         sample_json = self.__get_sample_json()