PYTHON-3396 Support the Azure VM-assigned Managed Identity for Automatic KMS Credentials (#1105)

Author: blink1073

.evergreen/config.yml

@@ -1072,6 +1072,50 @@ task_groups:
     tasks:
       - testgcpkms-task
 
+  - name: testazurekms_task_group
+    setup_group:
+    - func: fetch source
+    - func: prepare resources
+    - func: fix absolute paths
+    - func: make files executable
+    - command: shell.exec
+      params:
+        silent: true
+        shell: bash
+        script: |-
+          set -o errexit
+          ${PREPARE_SHELL}
+          echo '${testazurekms_publickey}' > /tmp/testazurekms_publickey
+          echo '${testazurekms_privatekey}' > /tmp/testazurekms_privatekey
+          # Set 600 permissions on private key file. Otherwise ssh / scp may error with permissions "are too open".
+          chmod 600 /tmp/testazurekms_privatekey
+          export AZUREKMS_CLIENTID="${testazurekms_clientid}"
+          export AZUREKMS_TENANTID="${testazurekms_tenantid}"
+          export AZUREKMS_SECRET="${testazurekms_secret}"
+          export AZUREKMS_DRIVERS_TOOLS="$DRIVERS_TOOLS"
+          export AZUREKMS_RESOURCEGROUP="${testazurekms_resourcegroup}"
+          export AZUREKMS_PUBLICKEYPATH="/tmp/testazurekms_publickey"
+          export AZUREKMS_PRIVATEKEYPATH="/tmp/testazurekms_privatekey"
+          export AZUREKMS_SCOPE="${testazurekms_scope}"
+          export AZUREKMS_VMNAME_PREFIX="PYTHON_DRIVER"
+          $DRIVERS_TOOLS/.evergreen/csfle/azurekms/create-and-setup-vm.sh
+    - command: expansions.update
+      params:
+        file: testazurekms-expansions.yml
+    teardown_group:
+    - command: shell.exec
+      params:
+        shell: bash
+        script: |-
+          ${PREPARE_SHELL}
+          export AZUREKMS_VMNAME=${AZUREKMS_VMNAME}
+          export AZUREKMS_RESOURCEGROUP=${testazurekms_resourcegroup}
+          $DRIVERS_TOOLS/.evergreen/csfle/azurekms/delete-vm.sh
+    setup_group_can_fail_task: true
+    setup_group_timeout_secs: 1800
+    tasks:
+    - testazurekms-task
+
 tasks:
     # Wildcard task. Do you need to find out what tools are available and where?
     # Throw it here, and execute this task on all buildvariants
@@ -1925,20 +1969,83 @@ tasks:
               export GCPKMS_PROJECT=${GCPKMS_PROJECT}
               export GCPKMS_ZONE=${GCPKMS_ZONE}
               export GCPKMS_INSTANCENAME=${GCPKMS_INSTANCENAME}
-              GCPKMS_CMD="SUCCESS=true ./.evergreen/run-mongodb-fle-gcp-auto.sh mongodb://localhost:27017" $DRIVERS_TOOLS/.evergreen/csfle/gcpkms/run-command.sh
+              GCPKMS_CMD="SUCCESS=true TEST_FLE_GCP_AUTO=1 LIBMONGOCRYPT_URL=https://s3.amazonaws.com/mciuploads/libmongocrypt/debian10/master/latest/libmongocrypt.tar.gz ./.evergreen/run-tests.sh" $DRIVERS_TOOLS/.evergreen/csfle/gcpkms/run-command.sh
 
     - name: "testgcpkms-fail-task"
       # testgcpkms-fail-task runs in a non-GCE environment.
       # It is expected to fail to obtain GCE credentials.
       commands:
+        - func: "bootstrap mongo-orchestration"
+          vars:
+            VERSION: "latest"
+            TOPOLOGY: "server"
         - command: shell.exec
           type: test
           params:
             working_dir: "src"
             shell: "bash"
             script: |
               ${PREPARE_SHELL}
-              SUCCESS=false ./.evergreen/run-mongodb-fle-gcp-auto.sh mongodb://localhost:27017
+              export LIBMONGOCRYPT_URL=https://s3.amazonaws.com/mciuploads/libmongocrypt/ubuntu1804-64/master/latest/libmongocrypt.tar.gz
+              SUCCESS=false TEST_FLE_GCP_AUTO=1 ./.evergreen/run-tests.sh
+
+    - name: testazurekms-task
+      commands:
+      - command: shell.exec
+        params:
+          shell: bash
+          script: |-
+            set -o errexit
+            ${PREPARE_SHELL}
+            cd src
+            echo "Copying files ... begin"
+            export AZUREKMS_RESOURCEGROUP=${testazurekms_resourcegroup}
+            export AZUREKMS_VMNAME=${AZUREKMS_VMNAME}
+            export AZUREKMS_PRIVATEKEYPATH=/tmp/testazurekms_privatekey
+            tar czf /tmp/mongo-python-driver.tgz .
+            AZUREKMS_SRC="/tmp/mongo-python-driver.tgz" \
+            AZUREKMS_DST="~/" \
+              $DRIVERS_TOOLS/.evergreen/csfle/azurekms/copy-file.sh
+            echo "Copying files ... end"
+            echo "Untarring file ... begin"
+            AZUREKMS_CMD="tar xf mongo-python-driver.tgz" \
+              $DRIVERS_TOOLS/.evergreen/csfle/azurekms/run-command.sh
+            echo "Untarring file ... end"
+      - command: shell.exec
+        type: test
+        params:
+          shell: bash
+          script: |-
+            set -o errexit
+            ${PREPARE_SHELL}
+            export AZUREKMS_RESOURCEGROUP=${testazurekms_resourcegroup}
+            export AZUREKMS_VMNAME=${AZUREKMS_VMNAME}
+            export AZUREKMS_PRIVATEKEYPATH=/tmp/testazurekms_privatekey
+            AZUREKMS_CMD="KEY_NAME='${testazurekms_keyname}' KEY_VAULT_ENDPOINT='${testazurekms_keyvaultendpoint}'  LIBMONGOCRYPT_URL=https://s3.amazonaws.com/mciuploads/libmongocrypt/debian10/master/latest/libmongocrypt.tar.gz SUCCESS=true TEST_FLE_AZURE_AUTO=1 ./.evergreen/run-tests.sh" \
+              $DRIVERS_TOOLS/.evergreen/csfle/azurekms/run-command.sh
+
+    - name: testazurekms-fail-task
+      commands:
+      - func: fetch source
+      - func: make files executable
+      - func: "bootstrap mongo-orchestration"
+        vars:
+          VERSION: "latest"
+          TOPOLOGY: "server"
+      - command: shell.exec
+        type: test
+        params:
+          shell: bash
+          script: |-
+            set -o errexit
+            ${PREPARE_SHELL}
+            cd src
+            PYTHON_BINARY=
+            KEY_NAME='${testazurekms_keyname}' \
+            KEY_VAULT_ENDPOINT='${testazurekms_keyvaultendpoint}' \
+            LIBMONGOCRYPT_URL=https://s3.amazonaws.com/mciuploads/libmongocrypt/ubuntu1804-64/master/latest/libmongocrypt.tar.gz \
+            SUCCESS=false TEST_FLE_AZURE_AUTO=1 \
+            ./.evergreen/run-tests.sh
 
 axes:
   # Choice of distro
@@ -2920,12 +3027,20 @@ buildvariants:
 - name: testgcpkms-variant
   display_name: "GCP KMS"
   run_on:
-    - debian11-small
+    - ubuntu1804-test
   tasks:
     - name: testgcpkms_task_group
       batchtime: 20160 # Use a batchtime of 14 days as suggested by the CSFLE test README
     - testgcpkms-fail-task
 
+- name: testazurekms-variant
+  display_name: "Azure KMS"
+  run_on: ubuntu1804-test
+  tasks:
+    - name: testazurekms_task_group
+      batchtime: 20160 # Use a batchtime of 14 days as suggested by the CSFLE test README
+    - testazurekms-fail-task
+
 - name: Release
   display_name: Release
   batchtime: 20160 # 14 days

.evergreen/run-mongodb-fle-gcp-auto.sh

@@ -101,7 +101,8 @@ if [ -n "$TEST_PYOPENSSL" ]; then
     python -m pip install --prefer-binary pyopenssl requests service_identity
 fi
 
-if [ -n "$TEST_ENCRYPTION" ]; then
+if [ -n "$TEST_ENCRYPTION" ] || [ -n "$TEST_FLE_AZURE_AUTO" ] || [ -n "$TEST_FLE_GCP_AUTO" ]; then
+
     createvirtualenv $PYTHON venv-encryption
     trap "deactivate; rm -rf venv-encryption" EXIT HUP
     PYTHON=python
@@ -146,7 +147,9 @@ if [ -n "$TEST_ENCRYPTION" ]; then
     python -c "import pymongocrypt; print('pymongocrypt version: '+pymongocrypt.__version__)"
     python -c "import pymongocrypt; print('libmongocrypt version: '+pymongocrypt.libmongocrypt_version())"
     # PATH is updated by PREPARE_SHELL for access to mongocryptd.
+fi
 
+if [ -n "$TEST_ENCRYPTION" ]; then
     # Need aws dependency for On-Demand KMS Credentials.
     python -m pip install '.[aws]'
 
@@ -171,6 +174,20 @@ if [ -n "$TEST_ENCRYPTION" ]; then
     TEST_ARGS="-s test.test_encryption"
 fi
 
+if [ -n "$TEST_FLE_AZURE_AUTO" ] || [ -n "$TEST_FLE_GCP_AUTO" ]; then
+    if [[ -z "$SUCCESS" ]]; then
+        echo "Must define SUCCESS"
+        exit 1
+    fi
+
+    if echo "$MONGODB_URI" | grep -q "@"; then
+      echo "MONGODB_URI unexpectedly contains user credentials in FLE test!";
+      exit 1
+    fi
+
+    TEST_ARGS="-s test.test_on_demand_csfle"
+fi
+
 if [ -n "$DATA_LAKE" ]; then
     TEST_ARGS="-s test.test_data_lake"
 fi

.evergreen/run-tests.sh

@@ -65,3 +65,47 @@ def test_02_success(self):
             codec_options=CodecOptions(),
         )
         self.client_encryption.create_data_key("gcp", self.master_key)
+
+
+class TestonDemandAzureCredentials(IntegrationTest):
+    @classmethod
+    @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed")
+    @client_context.require_version_min(4, 2, -1)
+    def setUpClass(cls):
+        super(TestonDemandAzureCredentials, cls).setUpClass()
+
+    def setUp(self):
+        super(TestonDemandAzureCredentials, self).setUp()
+        self.master_key = {
+            "keyVaultEndpoint": "https://keyvault-drivers-2411.vault.azure.net/keys/",
+            "keyName": "KEY-NAME",
+        }
+
+    @unittest.skipIf(not os.getenv("TEST_FLE_AZURE_AUTO"), "Not testing FLE Azure auto")
+    def test_01_failure(self):
+        if os.environ["SUCCESS"].lower() == "true":
+            self.skipTest("Expecting success")
+        self.client_encryption = ClientEncryption(
+            kms_providers={"azure": {}},
+            key_vault_namespace="keyvault.datakeys",
+            key_vault_client=client_context.client,
+            codec_options=CodecOptions(),
+        )
+        with self.assertRaises(EncryptionError):
+            self.client_encryption.create_data_key("azure", self.master_key)
+
+    @unittest.skipIf(not os.getenv("TEST_FLE_AZURE_AUTO"), "Not testing FLE Azure auto")
+    def test_02_success(self):
+        if os.environ["SUCCESS"].lower() == "false":
+            self.skipTest("Expecting failure")
+        self.client_encryption = ClientEncryption(
+            kms_providers={"azure": {}},
+            key_vault_namespace="keyvault.datakeys",
+            key_vault_client=client_context.client,
+            codec_options=CodecOptions(),
+        )
+        self.client_encryption.create_data_key("azure", self.master_key)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)