PYTHON-3467 OIDC: Automatic token acquisition for Azure Identity Provider

.evergreen/config.yml

@@ -98,7 +98,8 @@ functions:
             # If this was a patch build, doing a fresh clone would not actually test the patch
             cp -R ${PROJECT_DIRECTORY}/ $DRIVERS_TOOLS
           else
-            git clone https://github.com/mongodb-labs/drivers-evergreen-tools.git $DRIVERS_TOOLS
+            git clone --branch DRIVERS-2416-2 https://github.com/blink1073/drivers-evergreen-tools.git $DRIVERS_TOOLS
+            cat $DRIVERS_TOOLS/.evergreen/auth_oidc/azure/test.py
           fi
           echo "{ \"releases\": { \"default\": \"$MONGODB_BINARIES\" }}" > $MONGO_ORCHESTRATION_HOME/orchestration.config
 
@@ -1280,6 +1281,37 @@ task_groups:
     tasks:
     - testazurekms-task
 
+  - name: testazureoidc_task_group
+    setup_group:
+    - func: fetch source
+    - func: prepare resources
+    - func: fix absolute paths
+    - func: make files executable
+    - command: shell.exec
+      params:
+        shell: bash
+        script: |-
+          set -o errexit
+          ${PREPARE_SHELL}
+          export AZUREOIDC_CLIENTID="${testazureoidc_clientid}"
+          export AZUREOIDC_TENANTID="${testazureoic_tenantid}"
+          export AZUREOIDC_SECRET="${testazureoidc_secret}"
+          export AZUREOIDC_KEYVAULT=${testazureoidc_keyvault}
+          export AZUREOIDC_DRIVERS_TOOLS="$DRIVERS_TOOLS"
+          export AZUREOIDC_VMNAME_PREFIX="PYTHON_DRIVER"
+          $DRIVERS_TOOLS/.evergreen/auth_oidc/azure/create-and-setup-vm.sh
+    teardown_group:
+    - command: shell.exec
+      params:
+        shell: bash
+        script: |-
+          ${PREPARE_SHELL}
+          $DRIVERS_TOOLS/.evergreen/auth_oidc/azure/delete-vm.sh
+    setup_group_can_fail_task: true
+    setup_group_timeout_secs: 1800
+    tasks:
+    - oidc-auth-test-azure-latest
+
 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
@@ -2129,6 +2161,20 @@ tasks:
           vars:
             AWS_WEB_IDENTITY_TOKEN_FILE: /tmp/tokens/test1
 
+    - name: "oidc-auth-test-azure-latest"
+      commands:
+      - command: shell.exec
+        params:
+          shell: bash
+          script: |-
+            set -o errexit
+            ${PREPARE_SHELL}
+            cd src
+            export AZUREOIDC_DRIVERS_TAR_FILE=/tmp/mongo-python-driver.tgz
+            tar czf $AZUREOIDC_DRIVERS_TAR_FILE .
+            export AZUREOIDC_TEST_CMD="source ./env.sh && PROVIDER_NAME=azure ./.evergreen/run-mongodb-oidc-test.sh"
+            bash $DRIVERS_TOOLS/.evergreen/auth_oidc/azure/run-driver-test.sh
+
     - name: load-balancer-test
       commands:
         - func: "bootstrap mongo-orchestration"
@@ -3210,6 +3256,13 @@ buildvariants:
   tasks:
     - name: "oidc-auth-test-latest"
 
+- name: testazureoidc-variant
+  display_name: "Azure OIDC"
+  run_on: ubuntu2004-small
+  tasks:
+    - name: testazureoidc_task_group
+      batchtime: 20160 # Use a batchtime of 14 days as suggested by the CSFLE test README
+
 - matrix_name: "aws-auth-test"
   matrix_spec:
     platform: [ubuntu-18.04]

.evergreen/run-mongodb-oidc-test.sh

@@ -21,22 +21,26 @@ set +x
 shopt -s expand_aliases # needed for `urlencode` alias
 [ -s "${PROJECT_DIRECTORY}/prepare_mongodb_oidc.sh" ] && source "${PROJECT_DIRECTORY}/prepare_mongodb_oidc.sh"
 
-MONGODB_URI=${MONGODB_URI:-"mongodb://localhost"}
-MONGODB_URI_SINGLE="${MONGODB_URI}/?authMechanism=MONGODB-OIDC"
-MONGODB_URI_MULTIPLE="${MONGODB_URI}:27018/?authMechanism=MONGODB-OIDC&directConnection=true"
+PROVIDER_NAME=${PROVIDER_NAME:-"aws"}
+export MONGODB_URI=${MONGODB_URI:-"mongodb://localhost"}
 
-if [ -z "${OIDC_TOKEN_DIR}" ]; then
-    echo "Must specify OIDC_TOKEN_DIR"
-    exit 1
-fi
-
-export MONGODB_URI_SINGLE="$MONGODB_URI_SINGLE"
-export MONGODB_URI_MULTIPLE="$MONGODB_URI_MULTIPLE"
-export MONGODB_URI="$MONGODB_URI"
+if [ "$PROVIDER_NAME" = "aws" ]; then
+    export MONGODB_URI_SINGLE="${MONGODB_URI}/?authMechanism=MONGODB-OIDC"
+    export MONGODB_URI_MULTIPLE="${MONGODB_URI}:27018/?authMechanism=MONGODB-OIDC&directConnection=true"
 
-echo $MONGODB_URI_SINGLE
-echo $MONGODB_URI_MULTIPLE
-echo $MONGODB_URI
+    if [ -z "${OIDC_TOKEN_DIR}" ]; then
+        echo "Must specify OIDC_TOKEN_DIR"
+        exit 1
+    fi
+elif [ "$PROVIDER_NAME" = "azure" ]; then
+    if [ -z "${AZUREOIDC_CLIENTID}" ]; then
+        echo "Must specify an AZUREOIDC_CLIENTID"
+        exit 1
+    fi
+    MONGODB_URI="${MONGODB_URI}/?authMechanism=MONGODB-OIDC"
+    MONGODB_URI="${MONGODB_URI}&authMechanismProperties=PROVIDER_NAME:azure"
+    export MONGODB_URI="${MONGODB_URI},TOKEN_AUDIENCE:api%3A%2F%2F${AZUREOIDC_CLIENTID}"
+fi
 
 if [ "$ASSERT_NO_URI_CREDS" = "true" ]; then
     if echo "$MONGODB_URI" | grep -q "@"; then
@@ -47,13 +51,7 @@ fi
 
 # show test output
 set -x
-
-# Workaround macOS python 3.9 incompatibility with system virtualenv.
-if [ "$(uname -s)" = "Darwin" ]; then
-    VIRTUALENV="/Library/Frameworks/Python.framework/Versions/3.9/bin/python3 -m virtualenv"
-else
-    VIRTUALENV=$(command -v virtualenv)
-fi
+echo "MONGODB_URI=${MONGODB_URI}"
 
 authtest () {
     if [ "Windows_NT" = "$OS" ]; then
@@ -63,20 +61,27 @@ authtest () {
     echo "Running MONGODB-OIDC authentication tests with $PYTHON"
     $PYTHON --version
 
-    $VIRTUALENV -p $PYTHON --never-download venvoidc
+    $PYTHON -m venv venvoidc
     if [ "Windows_NT" = "$OS" ]; then
       . venvoidc/Scripts/activate
     else
       . venvoidc/bin/activate
     fi
     python -m pip install -U pip setuptools
-    python -m pip install '.[aws]'
-    python test/auth_aws/test_auth_oidc.py -v
+
+    if [ "$PROVIDER_NAME" = "aws" ]; then
+        python -m pip install '.[aws]'
+        python test/auth_oidc/test_auth_oidc.py -v
+    elif [ "$PROVIDER_NAME" = "azure" ]; then
+        python -m pip install .
+        python test/auth_oidc/test_auth_oidc_azure.py -v
+    fi
     deactivate
     rm -rf venvoidc
+    echo "MONGODB-OIDC authentication tests complete"
 }
 
-PYTHON=${PYTHON_BINARY:-}
+PYTHON=${PYTHON_BINARY:-$(which python3)}
 if [ -z "$PYTHON" ]; then
     echo "Cannot test without specifying PYTHON_BINARY"
     exit 1

pymongo/auth.py

@@ -157,6 +157,7 @@ def _build_credentials_tuple(
         request_token_callback = properties.get("request_token_callback")
         refresh_token_callback = properties.get("refresh_token_callback", None)
         provider_name = properties.get("PROVIDER_NAME", "")
+        token_audience = properties.get("TOKEN_AUDIENCE", "")
         default_allowed = [
             "*.mongodb.net",
             "*.mongodb-dev.net",
@@ -166,15 +167,22 @@ def _build_credentials_tuple(
             "::1",
         ]
         allowed_hosts = properties.get("allowed_hosts", default_allowed)
-        if not request_token_callback and provider_name != "aws":
+
+        # Handle providers.
+        providers = ["aws", "azure"]
+        if not request_token_callback and provider_name not in providers:
             raise ConfigurationError(
-                "authentication with MONGODB-OIDC requires providing an request_token_callback or a provider_name of 'aws'"
+                f"authentication with MONGODB-OIDC requires providing an request_token_callback or a provider_name that is one of {providers}"
             )
+        if provider_name == "azure" and not token_audience:
+            raise ConfigurationError("The azure provider requires a TOKEN_AUDIENCE value")
+
         oidc_props = _OIDCProperties(
             request_token_callback=request_token_callback,
             refresh_token_callback=refresh_token_callback,
             provider_name=provider_name,
             allowed_hosts=allowed_hosts,
+            token_audience=token_audience,
         )
         return MongoCredential(mech, "$external", user, passwd, oidc_props, None)
 

pymongo/auth_oidc.py

@@ -22,6 +22,7 @@
 import bson
 from bson.binary import Binary
 from bson.son import SON
+from pymongo.azure_helpers import _get_azure_token
 from pymongo.errors import ConfigurationError, OperationFailure
 from pymongo.helpers import _REAUTHENTICATION_REQUIRED_CODE
 
@@ -32,6 +33,7 @@ class _OIDCProperties:
     refresh_token_callback: Optional[Callable[..., Dict]]
     provider_name: Optional[str]
     allowed_hosts: List[str]
+    token_audience: Optional[str]
 
 
 """Mechanism properties for MONGODB-OIDC authentication."""
@@ -176,11 +178,20 @@ def get_current_token(self, use_callbacks=True):
     def auth_start_cmd(self, use_callbacks=True):
         properties = self.properties
 
-        # Handle aws provider credentials.
-        if properties.provider_name == "aws":
-            aws_identity_file = os.environ["AWS_WEB_IDENTITY_TOKEN_FILE"]
-            with open(aws_identity_file) as fid:
-                token = fid.read().strip()
+        # Handle provider credentials.
+        provider = properties.provider_name
+        if provider:
+            if provider == "aws":
+                aws_identity_file = os.environ["AWS_WEB_IDENTITY_TOKEN_FILE"]
+                with open(aws_identity_file) as fid:
+                    token = fid.read().strip()
+            elif provider == "azure":
+                assert self.properties.token_audience is not None
+                token = _get_azure_token(self.properties.token_audience)
+            else:
+                raise ConfigurationError(f"Unsupported provider {provider}")
+            self.token_gen_id += 1
+
             payload = {"jwt": token}
             cmd = SON(
                 [

pymongo/azure_helpers.py

@@ -0,0 +1,67 @@
+# Copyright 2023-present MongoDB, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Azure helpers."""
+import json
+from datetime import datetime, timedelta, timezone
+from typing import Dict, Tuple
+from urllib.request import Request, urlopen
+
+_CACHE: Dict[str, Tuple[str, datetime]] = {}
+_TOKEN_BUFFER_MINUTES = 5
+
+
+def _get_azure_token(resource: str, timeout: float = 5) -> str:
+    """Get an azure token for a given resource."""
+    cached_value = _CACHE.get(resource)
+    now_utc = datetime.now(timezone.utc)
+    if cached_value:
+        token, exp_utc = cached_value
+        buffer_seconds = _TOKEN_BUFFER_MINUTES * 60
+        if (exp_utc - now_utc).total_seconds() >= buffer_seconds:
+            return token
+        del _CACHE[resource]
+
+    url = "http://169.254.169.254/metadata/identity/oauth2/token"
+    url += "?api-version=2018-02-01"
+    url += f"&resource={resource}"
+    headers = {"Metadata": "true", "Accept": "application/json"}
+    request = Request(url, headers=headers)
+    try:
+        with urlopen(request, timeout=timeout) as response:
+            status = response.status
+            body = response.read().decode("utf8")
+    except Exception as e:
+        msg = "Failed to acquire IMDS access token: %s" % e
+        raise ValueError(msg)
+
+    if status != 200:
+        print(body)
+        msg = "Failed to acquire IMDS access token."
+        raise ValueError(msg)
+    try:
+        data = json.loads(body)
+    except Exception:
+        raise ValueError("Azure IMDS response must be in JSON format.")
+
+    for key in ["access_token", "expires_in"]:
+        if not data.get(key):
+            msg = "Azure IMDS response must contain %s, but was %s."
+            msg = msg % (key, body)
+            raise ValueError(msg)
+
+    token = data["access_token"]
+    exp_utc = now_utc + timedelta(seconds=int(data["expires_in"]))
+    _CACHE[resource] = (token, exp_utc)
+    return token

pymongo/common.py

@@ -421,6 +421,7 @@ def validate_read_preference_tags(name: str, value: Any) -> List[Dict[str, str]]
         "SERVICE_REALM",
         "AWS_SESSION_TOKEN",
         "PROVIDER_NAME",
+        "TOKEN_AUDIENCE",
     ]
 )