Skip to content

Commit 126f992

Browse files
pqlMCpqlMC
committed
fixed decorator response to be the same a the original but with decrypted body
fixed decorator tests
1 parent f0f7645 commit 126f992

File tree

2 files changed

+44
-35
lines changed

2 files changed

+44
-35
lines changed

‎client_encryption/api_encryption.py‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -29,10 +29,9 @@ def call_api_function(*args, **kwargs):
2929
kwargs["_preload_content"] = False
3030

3131
response = func(*args, **kwargs)
32+
response._body = self._decrypt_payload(response.getheaders(), response.data)
3233

33-
response_body = self._decrypt_payload(response.getheaders(), response.data)
34-
35-
return response_body
34+
return response
3635

3736
call_api_function.__fle__ = True
3837
return call_api_function
@@ -61,6 +60,7 @@ def _decrypt_payload(self, headers, body):
6160
"""Encryption enforcement based on configuration - decrypt using session key params from header or body"""
6261

6362
conf = self._encryption_conf
63+
params = None
6464

6565
if conf.use_http_headers:
6666
if conf.iv_field_name in headers and conf.encrypted_key_field_name in headers:
@@ -74,12 +74,12 @@ def _decrypt_payload(self, headers, body):
7474
del headers[conf.encryption_key_fingerprint_field_name]
7575

7676
params = SessionKeyParams(conf, encrypted_key, iv, oaep_digest_algo)
77-
payload = decrypt_payload(body, conf, params)
7877
else:
79-
# skip decryption if not iv nor key is in headers
80-
payload = body
81-
else:
82-
payload = decrypt_payload(body, conf)
78+
# skip decryption and return original body if not iv nor key is in headers
79+
return body
80+
81+
decrypted_body = decrypt_payload(body, conf, params)
82+
payload = json.dumps(decrypted_body, indent=4).encode('utf-8')
8383

8484
return payload
8585

‎tests/test_api_encryption.py‎

Lines changed: 36 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -62,14 +62,14 @@ def test_decrypt_payload_with_params_in_body(self):
6262

6363
test_headers = {"Content-Type": "application/json"}
6464

65-
decrypted = api_encryption._decrypt_payload(body={
65+
decrypted = json.loads(api_encryption._decrypt_payload(body={
6666
"encryptedData": {
6767
"iv": "uldLBySPY3VrznePihFYGQ==",
6868
"encryptedKey": "Jmh/bQPScUVFHSC9qinMGZ4lM7uetzUXcuMdEpC5g4C0Pb9HuaM3zC7K/509n7RTBZUPEzgsWtgi7m33nhpXsUo8WMcQkBIZlKn3ce+WRyZpZxcYtVoPqNn3benhcv7cq7yH1ktamUiZ5Dq7Ga+oQCaQEsOXtbGNS6vA5Bwa1pjbmMiRIbvlstInz8XTw8h/T0yLBLUJ0yYZmzmt+9i8qL8KFQ/PPDe5cXOCr1Aq2NTSixe5F2K/EI00q6D7QMpBDC7K6zDWgAOvINzifZ0DTkxVe4EE6F+FneDrcJsj+ZeIabrlRcfxtiFziH6unnXktta0sB1xcszIxXdMDbUcJA==",
6969
"encryptedValue": "KGfmdUWy89BwhQChzqZJ4w==",
7070
"oaepHashingAlgo": "SHA256"
7171
}
72-
}, headers=test_headers)
72+
}, headers=test_headers))
7373

7474
self.assertNotIn("encryptedData", decrypted)
7575
self.assertDictEqual({"data": {}}, decrypted)
@@ -112,11 +112,11 @@ def test_decrypt_payload_with_params_in_headers(self):
112112
}
113113

114114
api_encryption = to_test.ApiEncryption(self._json_config)
115-
decrypted = api_encryption._decrypt_payload(body={
115+
decrypted = json.loads(api_encryption._decrypt_payload(body={
116116
"encryptedData": {
117117
"encryptedValue": "KGfmdUWy89BwhQChzqZJ4w=="
118118
}
119-
}, headers=test_headers)
119+
}, headers=test_headers))
120120

121121
self.assertNotIn("encryptedData", decrypted)
122122
self.assertDictEqual({"data": {}}, decrypted)
@@ -170,9 +170,10 @@ def test_add_encryption_layer_post(self):
170170
}
171171
}, headers={"Content-Type": "application/json"})
172172

173-
self.assertIn("data", response)
174-
self.assertIn("secret", response["data"])
175-
self.assertEqual(secret2-secret1, response["data"]["secret"])
173+
self.assertIn("data", response.data)
174+
self.assertIn("secret", response.data["data"])
175+
self.assertEqual(secret2-secret1, response.data["data"]["secret"])
176+
self.assertDictEqual({"Content-Type": "application/json"}, response.getheaders())
176177

177178
def test_add_encryption_layer_delete(self):
178179
secret1 = 394
@@ -186,16 +187,18 @@ def test_add_encryption_layer_delete(self):
186187
}
187188
}, headers={"Content-Type": "application/json"})
188189

189-
self.assertEqual("OK", response)
190+
self.assertEqual("OK", response.data)
191+
self.assertDictEqual({"Content-Type": "application/json"}, response.getheaders())
190192

191193
def test_add_encryption_layer_get(self):
192194
test_client = MockApiClient()
193195
to_test.add_encryption_layer(test_client, self._json_config)
194196
response = MockService(test_client).do_something_get(headers={"Content-Type": "application/json"})
195197

196-
self.assertIn("data", response)
197-
self.assertIn("secret", response["data"])
198-
self.assertEqual([53, 84, 75], response["data"]["secret"])
198+
self.assertIn("data", response.data)
199+
self.assertIn("secret", response.data["data"])
200+
self.assertEqual([53, 84, 75], response.data["data"]["secret"])
201+
self.assertDictEqual({"Content-Type": "application/json"}, response.getheaders())
199202

200203
def test_add_header_encryption_layer_post_no_oaep_algo(self):
201204
self._set_header_params_config()
@@ -213,9 +216,10 @@ def test_add_header_encryption_layer_post_no_oaep_algo(self):
213216
"encryptedData": {}
214217
}, headers={"Content-Type": "application/json"})
215218

216-
self.assertIn("data", response)
217-
self.assertIn("secret", response["data"])
218-
self.assertEqual(secret2-secret1, response["data"]["secret"])
219+
self.assertIn("data", response.data)
220+
self.assertIn("secret", response.data["data"])
221+
self.assertEqual(secret2-secret1, response.data["data"]["secret"])
222+
self.assertDictEqual({"Content-Type": "application/json", "x-oaep-digest": "SHA256"}, response.getheaders())
219223

220224
def test_add_header_encryption_layer_post_no_cert_fingerprint(self):
221225
self._set_header_params_config()
@@ -233,9 +237,10 @@ def test_add_header_encryption_layer_post_no_cert_fingerprint(self):
233237
"encryptedData": {}
234238
}, headers={"Content-Type": "application/json"})
235239

236-
self.assertIn("data", response)
237-
self.assertIn("secret", response["data"])
238-
self.assertEqual(secret2-secret1, response["data"]["secret"])
240+
self.assertIn("data", response.data)
241+
self.assertIn("secret", response.data["data"])
242+
self.assertEqual(secret2-secret1, response.data["data"]["secret"])
243+
self.assertDictEqual({"Content-Type": "application/json"}, response.getheaders())
239244

240245
def test_add_header_encryption_layer_post_no_pubkey_fingerprint(self):
241246
self._set_header_params_config()
@@ -253,9 +258,10 @@ def test_add_header_encryption_layer_post_no_pubkey_fingerprint(self):
253258
"encryptedData": {}
254259
}, headers={"Content-Type": "application/json"})
255260

256-
self.assertIn("data", response)
257-
self.assertIn("secret", response["data"])
258-
self.assertEqual(secret2-secret1, response["data"]["secret"])
261+
self.assertIn("data", response.data)
262+
self.assertIn("secret", response.data["data"])
263+
self.assertEqual(secret2-secret1, response.data["data"]["secret"])
264+
self.assertDictEqual({"Content-Type": "application/json"}, response.getheaders())
259265

260266
def test_add_header_encryption_layer_no_iv(self):
261267
self._set_header_params_config()
@@ -288,9 +294,10 @@ def test_add_header_encryption_layer_post(self):
288294
"encryptedData": {}
289295
}, headers={"Content-Type": "application/json"})
290296

291-
self.assertIn("data", response)
292-
self.assertIn("secret", response["data"])
293-
self.assertEqual(secret2-secret1, response["data"]["secret"])
297+
self.assertIn("data", response.data)
298+
self.assertIn("secret", response.data["data"])
299+
self.assertEqual(secret2-secret1, response.data["data"]["secret"])
300+
self.assertDictEqual({"Content-Type": "application/json"}, response.getheaders())
294301

295302
def test_add_header_encryption_layer_delete(self):
296303
self._set_header_params_config()
@@ -307,7 +314,8 @@ def test_add_header_encryption_layer_delete(self):
307314
"encryptedData": {}
308315
}, headers={"Content-Type": "application/json"})
309316

310-
self.assertEqual("OK", response)
317+
self.assertEqual("OK", response.data)
318+
self.assertDictEqual({"Content-Type": "application/json"}, response.getheaders())
311319

312320
def test_add_header_encryption_layer_get(self):
313321
self._set_header_params_config()
@@ -316,9 +324,10 @@ def test_add_header_encryption_layer_get(self):
316324
to_test.add_encryption_layer(test_client, self._json_config)
317325
response = MockService(test_client).do_something_get_use_headers(headers={"Content-Type": "application/json"})
318326

319-
self.assertIn("data", response)
320-
self.assertIn("secret", response["data"])
321-
self.assertEqual([53, 84, 75], response["data"]["secret"])
327+
self.assertIn("data", response.data)
328+
self.assertIn("secret", response.data["data"])
329+
self.assertEqual([53, 84, 75], response.data["data"]["secret"])
330+
self.assertDictEqual({"Content-Type": "application/json"}, response.getheaders())
322331

323332
@patch('client_encryption.api_encryption.__oauth_warn')
324333
def test_add_encryption_layer_oauth_set(self, __oauth_warn):

0 commit comments

Comments
 (0)