[Cosmos] AAD authentication sync client (Azure#23604)

* working authentication to get database account

* working aad authentication for sync client with sample

* readme and changelog

* pylint and better comments on sample

* Update auth.py

* Revert "Update auth.py"

This reverts commit 721bbc7.

* Update auth.py

* Update auth.py

* changes from comments

* quick comment updates

* Update config.py

* Update access_cosmos_with_aad.py

* added sync policy to match async

* small changes

* aad tests for negative path and positive emulator path

* moved logic to be together for each part

* Milis comments

* Update cosmos_client.py

* Update dev_requirements.txt

* Update _auth_policy.py

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -3,12 +3,7 @@
 ### 4.3.0b4 (Unreleased)
 
 #### Features Added
-
-#### Breaking Changes
-
-#### Bugs Fixed
-
-#### Other Changes
+- Added support for AAD authentication for the sync client
 
 ### 4.3.0b3 (2022-03-10)
 

sdk/cosmos/azure-cosmos/README.md

@@ -76,6 +76,35 @@ KEY = os.environ['ACCOUNT_KEY']
 client = CosmosClient(URL, credential=KEY)
 ```
 
+### AAD Authentication
+
+You can also authenticate a client utilizing your service principal's AAD credentials and the azure identity package. 
+You can directly pass in the credentials information to ClientSecretCrednetial, or use the DefaultAzureCredential:
+```Python
+from azure.cosmos import CosmosClient
+from azure.identity import ClientSecretCredential, DefaultAzureCredential
+
+import os
+url = os.environ['ACCOUNT_URI']
+tenant_id = os.environ['TENANT_ID']
+client_id = os.environ['CLIENT_ID']
+client_secret = os.environ['CLIENT_SECRET']
+
+# Using ClientSecretCredential
+aad_credentials = ClientSecretCredential(
+    tenant_id=tenant_id,
+    client_id=client_id,
+    client_secret=client_secret)
+
+# Using DefaultAzureCredential (recommended)
+aad_credentials = DefaultAzureCredential()
+
+client = CosmosClient(url, aad_credentials)
+```
+Always ensure that the managed identity you use for AAD authentication has `readMetadata` permissions. <br>
+More information on how to set up AAD authentication: [Set up RBAC for AAD authentication](https://docs.microsoft.com/azure/cosmos-db/how-to-setup-rbac) <br>
+More information on allowed operations for AAD authenticated clients: [RBAC Permission Model](https://aka.ms/cosmos-native-rbac)
+
 ## Key concepts
 
 Once you've initialized a [CosmosClient][ref_cosmosclient], you can interact with the primary resource types in Cosmos DB:
@@ -125,7 +154,7 @@ Currently the features below are **not supported**. For alternatives options, ch
 * Change Feed: Processor
 * Change Feed: Read multiple partitions key values
 * Change Feed: Read specific time
-* Change Feed: Read from the beggining
+* Change Feed: Read from the beginning
 * Change Feed: Pull model
 * Cross-partition ORDER BY for mixed types
 
@@ -139,10 +168,6 @@ Currently the features below are **not supported**. For alternatives options, ch
 * Get the connection string
 * Get the minimum RU/s of a container
 
-### Security Limitations:
-
-* AAD support
-
 ## Workarounds
 
 ### Bulk processing Limitation Workaround
@@ -153,10 +178,6 @@ If you want to use Python SDK to perform bulk inserts to Cosmos DB, the best alt
 
 Typically, you can use [Azure Portal](https://portal.azure.com/), [Azure Cosmos DB Resource Provider REST API](https://docs.microsoft.com/rest/api/cosmos-db-resource-provider), [Azure CLI](https://docs.microsoft.com/cli/azure/azure-cli-reference-for-cosmos-db) or [PowerShell](https://docs.microsoft.com/azure/cosmos-db/manage-with-powershell) for the control plane unsupported limitations.
 
-### AAD Support Workaround
-
-A possible workaround is to use managed identities to [programmatically](https://docs.microsoft.com/azure/cosmos-db/managed-identity-based-authentication) get the keys.
-
 ## Boolean Data Type
 
 While the Python language [uses](https://docs.python.org/3/library/stdtypes.html?highlight=boolean#truth-value-testing) "True" and "False" for boolean types, Cosmos DB [accepts](https://docs.microsoft.com/azure/cosmos-db/sql-query-is-bool) "true" and "false" only. In other words, the Python language uses Boolean values with the first uppercase letter and all other lowercase letters, while Cosmos DB and its SQL language use only lowercase letters for those same Boolean values. How to deal with this challenge?

sdk/cosmos/azure-cosmos/azure/cosmos/_auth_policy.py

@@ -0,0 +1,166 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See LICENSE.txt in the project root for
+# license information.
+# -------------------------------------------------------------------------
+import time
+
+from typing import Any, Dict, Optional
+from azure.core.credentials import AccessToken
+from azure.core.pipeline import PipelineRequest, PipelineResponse
+from azure.core.pipeline.policies import HTTPPolicy
+from azure.cosmos import http_constants
+
+
+# pylint:disable=too-few-public-methods
+class _CosmosBearerTokenCredentialPolicyBase(object):
+    """Base class for a Bearer Token Credential Policy.
+
+    :param credential: The credential.
+    :type credential: ~azure.core.credentials.TokenCredential
+    :param str scopes: Lets you specify the type of access needed.
+    """
+
+    def __init__(self, credential, *scopes, **kwargs):  # pylint:disable=unused-argument
+        # type: (TokenCredential, *str, **Any) -> None
+        super(_CosmosBearerTokenCredentialPolicyBase, self).__init__()
+        self._scopes = scopes
+        self._credential = credential
+        self._token = None  # type: Optional[AccessToken]
+
+    @staticmethod
+    def _enforce_https(request):
+        # type: (PipelineRequest) -> None
+
+        # move 'enforce_https' from options to context so it persists
+        # across retries but isn't passed to a transport implementation
+        option = request.context.options.pop("enforce_https", None)
+
+        # True is the default setting; we needn't preserve an explicit opt in to the default behavior
+        if option is False:
+            request.context["enforce_https"] = option
+
+        enforce_https = request.context.get("enforce_https", True)
+        if enforce_https and not request.http_request.url.lower().startswith("https"):
+            raise ValueError(
+                "Bearer token authentication is not permitted for non-TLS protected (non-https) URLs."
+            )
+
+    @staticmethod
+    def _update_headers(headers, token):
+        # type: (Dict[str, str], str) -> None
+        """Updates the Authorization header with the bearer token.
+        This is the main method that differentiates this policy from core's BearerTokenCredentialPolicy and works
+        to properly sign the authorization header for Cosmos' REST API. For more information:
+        https://docs.microsoft.com/rest/api/cosmos-db/access-control-on-cosmosdb-resources#authorization-header
+
+        :param dict headers: The HTTP Request headers
+        :param str token: The OAuth token.
+        """
+        headers[http_constants.HttpHeaders.Authorization] = "type=aad&ver=1.0&sig={}".format(token)
+
+    @property
+    def _need_new_token(self):
+        # type: () -> bool
+        return not self._token or self._token.expires_on - time.time() < 300
+
+
+class CosmosBearerTokenCredentialPolicy(_CosmosBearerTokenCredentialPolicyBase, HTTPPolicy):
+    """Adds a bearer token Authorization header to requests.
+
+    :param credential: The credential.
+    :type credential: ~azure.core.TokenCredential
+    :param str scopes: Lets you specify the type of access needed.
+    :raises ValueError: If https_enforce does not match with endpoint being used.
+    """
+
+    def on_request(self, request):
+        # type: (PipelineRequest) -> None
+        """Called before the policy sends a request.
+
+        The base implementation authorizes the request with a bearer token.
+
+        :param ~azure.core.pipeline.PipelineRequest request: the request
+        """
+        self._enforce_https(request)
+
+        if self._token is None or self._need_new_token:
+            self._token = self._credential.get_token(*self._scopes)
+        self._update_headers(request.http_request.headers, self._token.token)
+
+    def authorize_request(self, request, *scopes, **kwargs):
+        # type: (PipelineRequest, *str, **Any) -> None
+        """Acquire a token from the credential and authorize the request with it.
+
+        Keyword arguments are passed to the credential's get_token method. The token will be cached and used to
+        authorize future requests.
+
+        :param ~azure.core.pipeline.PipelineRequest request: the request
+        :param str scopes: required scopes of authentication
+        """
+        self._token = self._credential.get_token(*scopes, **kwargs)
+        self._update_headers(request.http_request.headers, self._token.token)
+
+    def send(self, request):
+        # type: (PipelineRequest) -> PipelineResponse
+        """Authorize request with a bearer token and send it to the next policy
+
+        :param request: The pipeline request object
+        :type request: ~azure.core.pipeline.PipelineRequest
+        """
+        self.on_request(request)
+        try:
+            response = self.next.send(request)
+            self.on_response(request, response)
+        except Exception:  # pylint:disable=broad-except
+            self.on_exception(request)
+            raise
+        else:
+            if response.http_response.status_code == 401:
+                self._token = None  # any cached token is invalid
+                if "WWW-Authenticate" in response.http_response.headers:
+                    request_authorized = self.on_challenge(request, response)
+                    if request_authorized:
+                        try:
+                            response = self.next.send(request)
+                            self.on_response(request, response)
+                        except Exception:  # pylint:disable=broad-except
+                            self.on_exception(request)
+                            raise
+
+        return response
+
+    def on_challenge(self, request, response):
+        # type: (PipelineRequest, PipelineResponse) -> bool
+        """Authorize request according to an authentication challenge
+
+        This method is called when the resource provider responds 401 with a WWW-Authenticate header.
+
+        :param ~azure.core.pipeline.PipelineRequest request: the request which elicited an authentication challenge
+        :param ~azure.core.pipeline.PipelineResponse response: the resource provider's response
+        :returns: a bool indicating whether the policy should send the request
+        """
+        # pylint:disable=unused-argument,no-self-use
+        return False
+
+    def on_response(self, request, response):
+        # type: (PipelineRequest, PipelineResponse) -> None
+        """Executed after the request comes back from the next policy.
+
+        :param request: Request to be modified after returning from the policy.
+        :type request: ~azure.core.pipeline.PipelineRequest
+        :param response: Pipeline response object
+        :type response: ~azure.core.pipeline.PipelineResponse
+        """
+
+    def on_exception(self, request):
+        # type: (PipelineRequest) -> None
+        """Executed when an exception is raised while executing the next policy.
+
+        This method is executed inside the exception handler.
+
+        :param request: The Pipeline request object
+        :type request: ~azure.core.pipeline.PipelineRequest
+        """
+        # pylint: disable=no-self-use,unused-argument
+        return

sdk/cosmos/azure-cosmos/azure/cosmos/_base.py

@@ -30,6 +30,7 @@
 from typing import Dict, Any
 
 from urllib.parse import quote as urllib_quote
+from urllib.parse import urlsplit
 
 from azure.core import MatchConditions
 
@@ -663,6 +664,11 @@ def ParsePaths(paths):
     return tokens
 
 
+def create_scope_from_url(url):
+    parsed_url = urlsplit(url)
+    return parsed_url.scheme + "://" + parsed_url.hostname + "/.default"
+
+
 def validate_cache_staleness_value(max_integrated_cache_staleness):
     int(max_integrated_cache_staleness)  # Will throw error if data type cant be converted to int
     if max_integrated_cache_staleness <= 0:

sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py

@@ -58,6 +58,7 @@
 from . import _session
 from . import _utils
 from .partition_key import _Undefined, _Empty
+from ._auth_policy import CosmosBearerTokenCredentialPolicy
 
 ClassType = TypeVar("ClassType")
 
@@ -116,9 +117,11 @@ def __init__(
 
         self.master_key = None
         self.resource_tokens = None
+        self.aad_credentials = None
         if auth is not None:
             self.master_key = auth.get("masterKey")
             self.resource_tokens = auth.get("resourceTokens")
+            self.aad_credentials = auth.get("clientSecretCredential")
 
             if auth.get("permissionFeed"):
                 self.resource_tokens = {}
@@ -176,12 +179,18 @@ def __init__(
 
         self._user_agent = _utils.get_user_agent()
 
+        credentials_policy = None
+        if self.aad_credentials:
+            scopes = base.create_scope_from_url(self.url_connection)
+            credentials_policy = CosmosBearerTokenCredentialPolicy(self.aad_credentials, scopes)
+
         policies = [
             HeadersPolicy(**kwargs),
             ProxyPolicy(proxies=proxies),
             UserAgentPolicy(base_user_agent=self._user_agent, **kwargs),
             ContentDecodePolicy(),
             retry_policy,
+            credentials_policy,
             CustomHookPolicy(**kwargs),
             NetworkTraceLoggingPolicy(**kwargs),
             DistributedTracingPolicy(**kwargs),

sdk/cosmos/azure-cosmos/azure/cosmos/auth.py

@@ -26,12 +26,11 @@
 from hashlib import sha256
 import hmac
 import urllib.parse
-
 from . import http_constants
 
 
 def GetAuthorizationHeader(
-    cosmos_client_connection, verb, path, resource_id_or_fullname, is_name_based, resource_type, headers
+        cosmos_client_connection, verb, path, resource_id_or_fullname, is_name_based, resource_type, headers
 ):
     """Gets the authorization header.
 
@@ -51,18 +50,18 @@ def GetAuthorizationHeader(
         resource_id_or_fullname = resource_id_or_fullname.lower()
 
     if cosmos_client_connection.master_key:
-        return __GetAuthorizationTokenUsingMasterKey(
+        return __get_authorization_token_using_master_key(
             verb, resource_id_or_fullname, resource_type, headers, cosmos_client_connection.master_key
         )
     if cosmos_client_connection.resource_tokens:
-        return __GetAuthorizationTokenUsingResourceTokens(
+        return __get_authorization_token_using_resource_token(
             cosmos_client_connection.resource_tokens, path, resource_id_or_fullname
         )
 
     return None
 
 
-def __GetAuthorizationTokenUsingMasterKey(verb, resource_id_or_fullname, resource_type, headers, master_key):
+def __get_authorization_token_using_master_key(verb, resource_id_or_fullname, resource_type, headers, master_key):
     """Gets the authorization token using `master_key.
 
     :param str verb:
@@ -97,7 +96,7 @@ def __GetAuthorizationTokenUsingMasterKey(verb, resource_id_or_fullname, resourc
     return "type={type}&ver={ver}&sig={sig}".format(type=master_token, ver=token_version, sig=signature[:-1])
 
 
-def __GetAuthorizationTokenUsingResourceTokens(resource_tokens, path, resource_id_or_fullname):
+def __get_authorization_token_using_resource_token(resource_tokens, path, resource_id_or_fullname):
     """Get the authorization token using `resource_tokens`.
 
     :param dict resource_tokens:
@@ -138,7 +137,7 @@ def __GetAuthorizationTokenUsingResourceTokens(resource_tokens, path, resource_i
 
         # Get the last resource id or resource name from the path and get it's token from resource_tokens
         for i in range(len(path_parts), 1, -1):
-            segment = path_parts[i-1]
+            segment = path_parts[i - 1]
             sub_path = "/".join(path_parts[:i])
             if not segment in resource_types and sub_path in resource_tokens:
                 return resource_tokens[sub_path]

sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py

@@ -60,10 +60,13 @@ def _build_auth(credential):
         auth['resourceTokens'] = credential  # type: ignore
     elif hasattr(credential, '__iter__'):
         auth['permissionFeed'] = credential
+    elif hasattr(credential, 'get_token'):
+        auth['clientSecretCredential'] = credential
     else:
         raise TypeError(
-            "Unrecognized credential type. Please supply the master key as str, "
-            "or a dictionary or resource tokens, or a list of permissions.")
+            "Unrecognized credential type. Please supply the master key as a string "
+            "or a dictionary, or resource tokens, or a list of permissions, or any instance of a class implementing"
+            " TokenCredential (see azure.identity module for specific implementations such as ClientSecretCredential).")
     return auth
 
 

sdk/cosmos/azure-cosmos/azure/cosmos/http_constants.py

@@ -379,6 +379,7 @@ class SubStatusCodes(object):
     REDUNDANT_COLLECTION_PUT = 1009
     SHARED_THROUGHPUT_DATABASE_QUOTA_EXCEEDED = 1010
     SHARED_THROUGHPUT_OFFER_GROW_NOT_NEEDED = 1011
+    AAD_REQUEST_NOT_AUTHORIZED = 5300
 
     # 404: LSN in session token is higher
     READ_SESSION_NOTAVAILABLE = 1002

sdk/cosmos/azure-cosmos/dev_requirements.txt

@@ -1,3 +1,4 @@
 azure-core
+azure-identity
 -e ../../../tools/azure-sdk-tools
 -e ../../../tools/azure-devtools