PYTHON-2371 Add support for Azure and GCP KMS

prashantmital · prashantmital · commit cdbc3785c6ab · 2020-10-16T16:38:51.000-07:00
diff --git a/.evergreen/run-tests.sh b/.evergreen/run-tests.sh
@@ -152,7 +152,7 @@ if [ -n "$TEST_ENCRYPTION" ]; then
         exit 1
     fi
 
-    git clone --branch master git@github.com:mongodb/libmongocrypt.git libmongocrypt_git
+    git clone --branch PYTHON-2371/new-fle-keystores git@github.com:prashantmital/libmongocrypt.git libmongocrypt_git
     $PYTHON -m pip install --upgrade ./libmongocrypt_git/bindings/python
     # TODO: use a virtualenv
     trap "$PYTHON -m pip uninstall -y pymongocrypt" EXIT HUP
diff --git a/test/client-side-encryption/custom/azure-dek.json b/test/client-side-encryption/custom/azure-dek.json
@@ -0,0 +1,33 @@
+{
+    "_id": {
+        "$binary": {
+            "base64": "As3URE1jRcyHOPjaLWHOXA==",
+            "subType": "04"
+        }
+    },
+    "keyMaterial": {
+        "$binary": {
+            "base64": "df6fFLZqBsZSnQz2SnTYWNBtznIHktVSDMaidAdL7yVVgxBJQ0DyPZUR2HDQB4hdYym3w4C+VGqzcyTZNJOXn6nJzpGrGlIQMcjv93HE4sP2d245ShQCi1nTkLmMaXN63E2fzltOY3jW7ojf5Z4+r8kxmzyfymmSRgo0w8AF7lUWvFhnBYoE4tE322L31vtAK3Zj8pTPvw8/TcUdMSI9Y669IIzxbMy5yMPmdzpnb8nceUv6/CJoeiLhbt5GgaHqIAv7tHFOY8ZX8ztowMLa3GeAjd9clvzraDTqrfMFYco/kDKAW5iPQQ+Xuy1fP8tyFp0ZwaL/7Ed2sc819j8FTQ==",
+            "subType": "00"
+        }
+    },
+    "creationDate": {
+        "$date": {
+            "$numberLong": "1601573901680"
+        }
+    },
+    "updateDate": {
+        "$date": {
+            "$numberLong": "1601573901680"
+        }
+    },
+    "status": {
+        "$numberInt": "0"
+    },
+    "masterKey": {
+        "provider": "azure",
+        "keyVaultEndpoint": "key-vault-kevinalbs.vault.azure.net",
+        "keyName": "test-key"
+    }
+}
+
diff --git a/test/client-side-encryption/custom/azure-gcp-schema.json b/test/client-side-encryption/custom/azure-gcp-schema.json
@@ -0,0 +1,32 @@
+{
+    "db.coll": {
+        "bsonType": "object",
+        "properties": {
+            "secret_azure": {
+                "encrypt": {
+                    "keyId": [{
+                        "$binary": {
+                            "base64": "As3URE1jRcyHOPjaLWHOXA==",
+                            "subType": "04"
+                        }
+                    }],
+                    "algorithm": "AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic",
+                    "bsonType": "string"
+                }
+            },
+            "secret_gcp": {
+                "encrypt": {
+                    "keyId": [{
+                        "$binary": {
+                            "base64": "osU8SLxJRHONbl8Oh5o+eg==",
+                            "subType": "04"
+                        }
+                    }],
+                    "algorithm": "AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic",
+                    "bsonType": "string"
+                }
+            }
+        }
+    }
+}
+
diff --git a/test/client-side-encryption/custom/gcp-dek.json b/test/client-side-encryption/custom/gcp-dek.json
@@ -0,0 +1,35 @@
+{
+    "_id": {
+        "$binary": {
+            "base64": "osU8SLxJRHONbl8Oh5o+eg==",
+            "subType": "04"
+        }
+    },
+    "keyMaterial": {
+        "$binary": {
+            "base64": "CiQAg4LDql74hjYPZ957Z7YpCrD6yTVVXKegflJDstQ/xngTyx0SiQEAkWNo/fjPj6jMNSvEop07/29Fu72QHFDRYM3e/KFHfnMQjKzfxb1yX1dC6MbO5FZG/UNBkXlJgPqbHNVuizea3QC24kV5iOiEb4nTM7+RW+8TfVb6QerWWe6MjC+kNpj4LMVcc1lFfVDeGgpJLyMLNGitrjR16qH8qQTNbGNy0toTL69JUmgS8Q==",
+            "subType": "00"
+        }
+    },
+    "creationDate": {
+        "$date": {
+            "$numberLong": "1601574333107"
+        }
+    },
+    "updateDate": {
+        "$date": {
+            "$numberLong": "1601574333107"
+        }
+    },
+    "status": {
+        "$numberInt": "0"
+    },
+    "masterKey": {
+        "provider": "gcp",
+        "projectId": "csfle-poc",
+        "location": "global",
+        "keyRing": "test",
+        "keyName": "quickstart"
+    }
+}
+
diff --git a/test/test_encryption.py b/test/test_encryption.py
@@ -20,6 +20,7 @@
 import traceback
 import socket
 import sys
+import textwrap
 import uuid
 
 sys.path[0:0] = [""]
@@ -52,6 +53,7 @@
 from test.utils import (TestCreator,
                         camel_to_snake_args,
                         OvertCommandListener,
+                        WhiteListEventListener,
                         rs_or_single_client,
                         wait_until)
 from test.utils_spec_runner import SpecRunner
@@ -1105,5 +1107,141 @@ def test_05_endpoint_invalid_host(self):
                 'aws', master_key=master_key)
 
 
+class AzureGCPEncryptionTestMixin(object):
+    KMS_PROVIDER_MAP = None
+    KEYVAULT_DB = 'keyvault'
+    KEYVAULT_COLL = 'datakeys'
+
+    def read_extjson(self, filename):
+        with open(os.path.join(BASE, 'custom', filename), 'r') as fp:
+            return json_util.loads(fp.read(), json_options=JSON_OPTS)
+
+    def setup_keyvault(self, DEK):
+        keyvault = self.client.get_database(self.KEYVAULT_DB).get_collection(
+            self.KEYVAULT_COLL, codec_options=OPTS,
+            write_concern=WriteConcern('majority'))
+        keyvault.drop()
+        keyvault.insert_one(DEK)
+
+    def _test_explicit(self, expectation):
+        client_encryption = ClientEncryption(
+            self.KMS_PROVIDER_MAP,
+            '.'.join([self.KEYVAULT_DB, self.KEYVAULT_COLL]),
+            client_context.client,
+            OPTS)
+
+        ciphertext = client_encryption.encrypt(
+            'test',
+            algorithm=Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic,
+            key_id=Binary.from_uuid(self.DEK['_id'], STANDARD))
+
+        self.assertEqual(bytes(ciphertext), base64.b64decode(expectation))
+        self.assertEqual(client_encryption.decrypt(ciphertext), 'test')
+
+    def _test_automatic(self, expectation_extjson, payload):
+        encrypted_db = "db"
+        encrypted_coll = "coll"
+
+        encryption_opts = AutoEncryptionOpts(
+            self.KMS_PROVIDER_MAP,
+            '.'.join([self.KEYVAULT_DB, self.KEYVAULT_COLL]),
+            schema_map=self.SCHEMA_MAP)
+
+        insert_listener = WhiteListEventListener('insert')
+        client = rs_or_single_client(
+            auto_encryption_opts=encryption_opts,
+            event_listeners=[insert_listener])
+
+        coll = client.get_database(encrypted_db).get_collection(
+            encrypted_coll, codec_options=OPTS,
+            write_concern=WriteConcern("majority"))
+        coll.drop()
+
+        expected_document = json_util.loads(
+            expectation_extjson, json_options=JSON_OPTS)
+
+        coll.insert_one(payload)
+        event = insert_listener.results['started'][0]
+        inserted_doc = event.command['documents'][0]
+
+        for key, value in expected_document.items():
+            self.assertEqual(value, inserted_doc[key])
+
+        output_doc = coll.find_one({})
+        for key, value in payload.items():
+            self.assertEqual(output_doc[key], value)
+
+
+AZURE_CREDS = {
+    'tenantId': os.environ.get('FLE_AZURE_TENANTID', ''),
+    'clientId': os.environ.get('FLE_AZURE_CLIENTID', ''),
+    'clientSecret': os.environ.get('FLE_AZURE_CLIENTSECRET', '')}
+
+
+class TestAzureEncryption(AzureGCPEncryptionTestMixin,
+                          EncryptionIntegrationTest):
+    @classmethod
+    @unittest.skipUnless(_HAVE_PYMONGOCRYPT, 'pymongocrypt is not installed')
+    @unittest.skipUnless(any(AZURE_CREDS.values()),
+                         'Azure environment credentials are not set')
+    def setUpClass(cls):
+        cls.KMS_PROVIDER_MAP = {'azure': AZURE_CREDS}
+        cls.DEK = json_data(BASE, 'custom', 'azure-dek.json')
+        cls.SCHEMA_MAP = json_data(BASE, 'custom', 'azure-gcp-schema.json')
+        super(TestAzureEncryption, cls).setUpClass()
+
+    def setUp(self):
+        self.setup_keyvault(self.DEK)
+
+    def test_explicit(self):
+        return self._test_explicit(
+            'AQLN1ERNY0XMhzj42i1hzlwC8/OSU9bHfaQRmmRF5l7d5ZpqJX13qF5zSyExo8N9c1b6uS/LoKrHNzcEMKNrkpi3jf2HiShTFRF0xi8AOD9yfw==')
+
+    def test_automatic(self):
+        expected_document_extjson = textwrap.dedent(""" 
+        {"secret_azure": {
+            "$binary": {
+                "base64": "AQLN1ERNY0XMhzj42i1hzlwC8/OSU9bHfaQRmmRF5l7d5ZpqJX13qF5zSyExo8N9c1b6uS/LoKrHNzcEMKNrkpi3jf2HiShTFRF0xi8AOD9yfw==",
+                "subType": "06"}
+        }}""")
+        return self._test_automatic(
+            expected_document_extjson, {"secret_azure": "test"})
+
+
+GCP_CREDS = {
+    'email': os.environ.get('FLE_GCP_EMAIL', ''),
+    'privateKey': os.environ.get('FLE_GCP_PRIVATEKEY', '')}
+
+
+class TestGCPEncryption(AzureGCPEncryptionTestMixin,
+                        EncryptionIntegrationTest):
+    @classmethod
+    @unittest.skipUnless(_HAVE_PYMONGOCRYPT, 'pymongocrypt is not installed')
+    @unittest.skipUnless(any(GCP_CREDS.values()),
+                         'GCP environment credentials are not set')
+    def setUpClass(cls):
+        cls.KMS_PROVIDER_MAP = {'gcp': GCP_CREDS}
+        cls.DEK = json_data(BASE, 'custom', 'gcp-dek.json')
+        cls.SCHEMA_MAP = json_data(BASE, 'custom', 'azure-gcp-schema.json')
+        super(TestGCPEncryption, cls).setUpClass()
+
+    def setUp(self):
+        self.setup_keyvault(self.DEK)
+
+    def test_explicit(self):
+        return self._test_explicit(
+            'AaLFPEi8SURzjW5fDoeaPnoCGcOFAmFOPpn5584VPJJ8iXIgml3YDxMRZD9IWv5otyoft8fBzL1LsDEp0lTeB32cV1gOj0IYeAKHhGIleuHZtA==')
+
+    def test_automatic(self):
+        expected_document_extjson = textwrap.dedent(""" 
+        {"secret_gcp": {
+            "$binary": {
+                "base64": "AaLFPEi8SURzjW5fDoeaPnoCGcOFAmFOPpn5584VPJJ8iXIgml3YDxMRZD9IWv5otyoft8fBzL1LsDEp0lTeB32cV1gOj0IYeAKHhGIleuHZtA==",
+                "subType": "06"}
+        }}""")
+        return self._test_automatic(
+            expected_document_extjson, {"secret_gcp": "test"})
+
+
 if __name__ == "__main__":
     unittest.main()

-Original file line number
+Diff line change
         exit 1
     fi
 -    git clone --branch master [email protected]:mongodb/libmongocrypt.git libmongocrypt_git
 +    git clone --branch PYTHON-2371/new-fle-keystores [email protected]:prashantmital/libmongocrypt.git libmongocrypt_git
     $PYTHON -m pip install --upgrade ./libmongocrypt_git/bindings/python
     # TODO: use a virtualenv
     trap "$PYTHON -m pip uninstall -y pymongocrypt" EXIT HUP