Azure · simorenoh · Apr 6, 2022 · Mar 18, 2022 · Mar 22, 2022 · Mar 22, 2022
@@ -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)
 

@@ -76,6 +76,28 @@ 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:
+```Python
+from azure.cosmos import CosmosClient
+from azure.identity import ClientSecretCredential
+
+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']
+
+AAD_CREDENTIALS = ClientSecretCredential(
+    tenant_id=TENANT_ID,
+    client_id=CLIENT_ID,
+    client_secret=CLIENT_SECRET)
+
+client = CosmosClient(URL, AAD_CREDENTIALS)
+```
+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 +147,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 +161,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 +171,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?

@@ -663,6 +663,11 @@ def ParsePaths(paths):
     return tokens
 
 
+def create_scope_from_url(url):
+    # Need a better way to do this, regex?
+    return url.replace(":443", "") + ".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:

@@ -58,6 +58,7 @@
 from . import _session
 from . import _utils
 from .partition_key import _Undefined, _Empty
+from .auth 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 = {}
@@ -188,6 +191,10 @@ def __init__(
             HttpLoggingPolicy(**kwargs),
         ]
 
+        if self.aad_credentials:
+            scopes = base.create_scope_from_url(self.url_connection)
+            policies.append(CosmosBearerTokenCredentialPolicy(self.aad_credentials, scopes))
+
         transport = kwargs.pop("transport", None)
         self.pipeline_client = PipelineClient(base_url=url_connection, transport=transport, policies=policies)
 

@@ -26,12 +26,17 @@
 from hashlib import sha256
 import hmac
 import urllib.parse
-
+import time
+from typing import Any, Dict, Optional, ClassVar
+from azure.core.exceptions import ServiceRequestError
+from azure.core.credentials import AccessToken
+from azure.core.pipeline import PipelineRequest, PipelineResponse
 from . import http_constants
 
+TokenCredential = ClassVar
 
 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 +56,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 +102,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,9 +143,166 @@ 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]
 
     return None
+
+
+# pylint:disable=too-few-public-methods
+class _CosmosBearerTokenCredentialPolicyBase(object):
+    """Base class for a Cosmos Bearer Token Credential Policy.
+    Based on ~azure.core.pipeline.policies.BearerTokenCredentialPolicy.
+    Takes care of updating token as needed and updating the request headers.
+
+    :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 ServiceRequestError(
+                "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 Cosmos AAD bearer token.
+
+        :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):
+    """Adds a Cosmos 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: :class:`~azure.core.exceptions.ServiceRequestError`
+    """
+
+    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
+            handled = self.on_exception(request)
+            if not handled:
+                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
+                            handled = self.on_exception(request)
+                            if not handled:
+                                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) -> bool
+        """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
+        :return: False by default, override with True to stop the exception.
+        :rtype: bool
+        """
+        # pylint: disable=no-self-use,unused-argument
+        return False
@@ -60,10 +60,12 @@ 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 a ClientSecretCredential.")
     return auth
 
 

@@ -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