|
20 | 20 | import traceback |
21 | 21 | import socket |
22 | 22 | import sys |
| 23 | +import textwrap |
23 | 24 | import uuid |
24 | 25 |
|
25 | 26 | sys.path[0:0] = [""] |
|
30 | 31 | STANDARD, |
31 | 32 | UUID_SUBTYPE) |
32 | 33 | from bson.codec_options import CodecOptions |
| 34 | +from bson.py3compat import _unicode |
33 | 35 | from bson.errors import BSONError |
34 | 36 | from bson.json_util import JSONOptions |
35 | 37 | from bson.son import SON |
|
52 | 54 | from test.utils import (TestCreator, |
53 | 55 | camel_to_snake_args, |
54 | 56 | OvertCommandListener, |
| 57 | + WhiteListEventListener, |
55 | 58 | rs_or_single_client, |
56 | 59 | wait_until) |
57 | 60 | from test.utils_spec_runner import SpecRunner |
@@ -1105,5 +1108,132 @@ def test_05_endpoint_invalid_host(self): |
1105 | 1108 | 'aws', master_key=master_key) |
1106 | 1109 |
|
1107 | 1110 |
|
| 1111 | +class AzureGCPEncryptionTestMixin(object): |
| 1112 | + DEK = None |
| 1113 | + KMS_PROVIDER_MAP = None |
| 1114 | + KEYVAULT_DB = 'keyvault' |
| 1115 | + KEYVAULT_COLL = 'datakeys' |
| 1116 | + |
| 1117 | + def setUp(self): |
| 1118 | + keyvault = self.client.get_database( |
| 1119 | + self.KEYVAULT_DB).get_collection( |
| 1120 | + self.KEYVAULT_COLL) |
| 1121 | + create_key_vault(keyvault, self.DEK) |
| 1122 | + |
| 1123 | + def _test_explicit(self, expectation): |
| 1124 | + client_encryption = ClientEncryption( |
| 1125 | + self.KMS_PROVIDER_MAP, |
| 1126 | + '.'.join([self.KEYVAULT_DB, self.KEYVAULT_COLL]), |
| 1127 | + client_context.client, |
| 1128 | + OPTS) |
| 1129 | + self.addCleanup(client_encryption.close) |
| 1130 | + |
| 1131 | + ciphertext = client_encryption.encrypt( |
| 1132 | + 'test', |
| 1133 | + algorithm=Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic, |
| 1134 | + key_id=Binary.from_uuid(self.DEK['_id'], STANDARD)) |
| 1135 | + |
| 1136 | + self.assertEqual(bytes(ciphertext), base64.b64decode(expectation)) |
| 1137 | + self.assertEqual(client_encryption.decrypt(ciphertext), 'test') |
| 1138 | + |
| 1139 | + def _test_automatic(self, expectation_extjson, payload): |
| 1140 | + encrypted_db = "db" |
| 1141 | + encrypted_coll = "coll" |
| 1142 | + keyvault_namespace = '.'.join([self.KEYVAULT_DB, self.KEYVAULT_COLL]) |
| 1143 | + |
| 1144 | + encryption_opts = AutoEncryptionOpts( |
| 1145 | + self.KMS_PROVIDER_MAP, |
| 1146 | + keyvault_namespace, |
| 1147 | + schema_map=self.SCHEMA_MAP) |
| 1148 | + |
| 1149 | + insert_listener = WhiteListEventListener('insert') |
| 1150 | + client = rs_or_single_client( |
| 1151 | + auto_encryption_opts=encryption_opts, |
| 1152 | + event_listeners=[insert_listener]) |
| 1153 | + self.addCleanup(client.close) |
| 1154 | + |
| 1155 | + coll = client.get_database(encrypted_db).get_collection( |
| 1156 | + encrypted_coll, codec_options=OPTS, |
| 1157 | + write_concern=WriteConcern("majority")) |
| 1158 | + coll.drop() |
| 1159 | + |
| 1160 | + expected_document = json_util.loads( |
| 1161 | + expectation_extjson, json_options=JSON_OPTS) |
| 1162 | + |
| 1163 | + coll.insert_one(payload) |
| 1164 | + event = insert_listener.results['started'][0] |
| 1165 | + inserted_doc = event.command['documents'][0] |
| 1166 | + |
| 1167 | + for key, value in expected_document.items(): |
| 1168 | + self.assertEqual(value, inserted_doc[key]) |
| 1169 | + |
| 1170 | + output_doc = coll.find_one({}) |
| 1171 | + for key, value in payload.items(): |
| 1172 | + self.assertEqual(output_doc[key], value) |
| 1173 | + |
| 1174 | + |
| 1175 | +AZURE_CREDS = { |
| 1176 | + 'tenantId': os.environ.get('FLE_AZURE_TENANTID', ''), |
| 1177 | + 'clientId': os.environ.get('FLE_AZURE_CLIENTID', ''), |
| 1178 | + 'clientSecret': os.environ.get('FLE_AZURE_CLIENTSECRET', '')} |
| 1179 | + |
| 1180 | + |
| 1181 | +class TestAzureEncryption(AzureGCPEncryptionTestMixin, |
| 1182 | + EncryptionIntegrationTest): |
| 1183 | + @classmethod |
| 1184 | + @unittest.skipUnless(any(AZURE_CREDS.values()), |
| 1185 | + 'Azure environment credentials are not set') |
| 1186 | + def setUpClass(cls): |
| 1187 | + cls.KMS_PROVIDER_MAP = {'azure': AZURE_CREDS} |
| 1188 | + cls.DEK = json_data(BASE, 'custom', 'azure-dek.json') |
| 1189 | + cls.SCHEMA_MAP = json_data(BASE, 'custom', 'azure-gcp-schema.json') |
| 1190 | + super(TestAzureEncryption, cls).setUpClass() |
| 1191 | + |
| 1192 | + def test_explicit(self): |
| 1193 | + return self._test_explicit( |
| 1194 | + 'AQLN1ERNY0XMhzj42i1hzlwC8/OSU9bHfaQRmmRF5l7d5ZpqJX13qF5zSyExo8N9c1b6uS/LoKrHNzcEMKNrkpi3jf2HiShTFRF0xi8AOD9yfw==') |
| 1195 | + |
| 1196 | + def test_automatic(self): |
| 1197 | + expected_document_extjson = textwrap.dedent(""" |
| 1198 | + {"secret_azure": { |
| 1199 | + "$binary": { |
| 1200 | + "base64": "AQLN1ERNY0XMhzj42i1hzlwC8/OSU9bHfaQRmmRF5l7d5ZpqJX13qF5zSyExo8N9c1b6uS/LoKrHNzcEMKNrkpi3jf2HiShTFRF0xi8AOD9yfw==", |
| 1201 | + "subType": "06"} |
| 1202 | + }}""") |
| 1203 | + return self._test_automatic( |
| 1204 | + expected_document_extjson, {"secret_azure": "test"}) |
| 1205 | + |
| 1206 | + |
| 1207 | +GCP_CREDS = { |
| 1208 | + 'email': os.environ.get('FLE_GCP_EMAIL', ''), |
| 1209 | + 'privateKey': _unicode(os.environ.get('FLE_GCP_PRIVATEKEY', ''))} |
| 1210 | + |
| 1211 | + |
| 1212 | +class TestGCPEncryption(AzureGCPEncryptionTestMixin, |
| 1213 | + EncryptionIntegrationTest): |
| 1214 | + @classmethod |
| 1215 | + @unittest.skipUnless(any(GCP_CREDS.values()), |
| 1216 | + 'GCP environment credentials are not set') |
| 1217 | + def setUpClass(cls): |
| 1218 | + cls.KMS_PROVIDER_MAP = {'gcp': GCP_CREDS} |
| 1219 | + cls.DEK = json_data(BASE, 'custom', 'gcp-dek.json') |
| 1220 | + cls.SCHEMA_MAP = json_data(BASE, 'custom', 'azure-gcp-schema.json') |
| 1221 | + super(TestGCPEncryption, cls).setUpClass() |
| 1222 | + |
| 1223 | + def test_explicit(self): |
| 1224 | + return self._test_explicit( |
| 1225 | + 'AaLFPEi8SURzjW5fDoeaPnoCGcOFAmFOPpn5584VPJJ8iXIgml3YDxMRZD9IWv5otyoft8fBzL1LsDEp0lTeB32cV1gOj0IYeAKHhGIleuHZtA==') |
| 1226 | + |
| 1227 | + def test_automatic(self): |
| 1228 | + expected_document_extjson = textwrap.dedent(""" |
| 1229 | + {"secret_gcp": { |
| 1230 | + "$binary": { |
| 1231 | + "base64": "AaLFPEi8SURzjW5fDoeaPnoCGcOFAmFOPpn5584VPJJ8iXIgml3YDxMRZD9IWv5otyoft8fBzL1LsDEp0lTeB32cV1gOj0IYeAKHhGIleuHZtA==", |
| 1232 | + "subType": "06"} |
| 1233 | + }}""") |
| 1234 | + return self._test_automatic( |
| 1235 | + expected_document_extjson, {"secret_gcp": "test"}) |
| 1236 | + |
| 1237 | + |
1108 | 1238 | if __name__ == "__main__": |
1109 | 1239 | unittest.main() |
0 commit comments