[Key Vault] Add support for multi-tenant authentication

sdk/keyvault/azure-keyvault-administration/CHANGELOG.md

@@ -3,12 +3,15 @@
 ## 4.1.0b2 (Unreleased)
 
 ### Features Added
+- Added support for multi-tenant authentication when using `azure-identity` 1.7.1 or newer
+  ([#20698](https://github.com/Azure/azure-sdk-for-python/issues/20698))
 
 ### Breaking Changes
 
 ### Bugs Fixed
 
 ### Other Changes
+- Updated minimum `azure-core` version to 1.15.0
 
 ## 4.1.0b1 (2021-09-09)
 

sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/_internal/__init__.py

@@ -6,15 +6,14 @@
 
 from six.moves.urllib_parse import urlparse
 
-from .challenge_auth_policy import ChallengeAuthPolicy, ChallengeAuthPolicyBase
+from .challenge_auth_policy import ChallengeAuthPolicy
 from .client_base import KeyVaultClientBase
 from .http_challenge import HttpChallenge
 from . import http_challenge_cache as HttpChallengeCache
 
 
 __all__ = [
     "ChallengeAuthPolicy",
-    "ChallengeAuthPolicyBase",
     "HttpChallenge",
     "HttpChallengeCache",
     "KeyVaultClientBase",

sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/_internal/async_challenge_auth_policy.py

@@ -13,67 +13,67 @@
 requirements can change. For example, a vault may move to a new tenant. In such a case the policy will attempt the
 protocol again.
 """
+
+import time
 from typing import TYPE_CHECKING
 
-from azure.core.pipeline.policies import AsyncHTTPPolicy
+from azure.core.pipeline.policies import AsyncBearerTokenCredentialPolicy
 
-from . import HttpChallengeCache
-from .challenge_auth_policy import _enforce_tls, _get_challenge_request, _update_challenge, ChallengeAuthPolicyBase
+from . import http_challenge_cache as ChallengeCache
+from .challenge_auth_policy import _enforce_tls, _get_challenge_scope, _update_challenge
 
 if TYPE_CHECKING:
-    from typing import Any
+    from typing import Any, Optional
+    from azure.core.credentials import AccessToken
     from azure.core.credentials_async import AsyncTokenCredential
-    from azure.core.pipeline import PipelineRequest
-    from azure.core.pipeline.transport import AsyncHttpResponse
-    from . import HttpChallenge
+    from azure.core.pipeline import PipelineRequest, PipelineResponse
 
 
-class AsyncChallengeAuthPolicy(ChallengeAuthPolicyBase, AsyncHTTPPolicy):
+class AsyncChallengeAuthPolicy(AsyncBearerTokenCredentialPolicy):
     """policy for handling HTTP authentication challenges"""
 
-    def __init__(self, credential: "AsyncTokenCredential", **kwargs: "Any") -> None:
+    def __init__(self, credential: "AsyncTokenCredential", *scopes: str, **kwargs: "Any") -> None:
+        super().__init__(credential, *scopes, **kwargs)
         self._credential = credential
-        super(AsyncChallengeAuthPolicy, self).__init__(**kwargs)
+        self._token = None  # type: Optional[AccessToken]
 
-    async def send(self, request: "PipelineRequest") -> "AsyncHttpResponse":
+    async def on_request(self, request: "PipelineRequest") -> None:
         _enforce_tls(request)
-
-        challenge = HttpChallengeCache.get_challenge_for_url(request.http_request.url)
-        if not challenge:
-            challenge_request = _get_challenge_request(request)
-            challenger = await self.next.send(challenge_request)
-            try:
-                challenge = _update_challenge(request, challenger)
-            except ValueError:
-                # didn't receive the expected challenge -> nothing more this policy can do
-                return challenger
-
-        await self._handle_challenge(request, challenge)
-        response = await self.next.send(request)
-
-        if response.http_response.status_code == 401:
-            # any cached token must be invalid
-            self._token = None
-
-            # cached challenge could be outdated; maybe this response has a new one?
-            try:
-                challenge = _update_challenge(request, response)
-            except ValueError:
-                # 401 with no legible challenge -> nothing more this policy can do
-                return response
-
-            await self._handle_challenge(request, challenge)
-            response = await self.next.send(request)
-
-        return response
-
-    async def _handle_challenge(self, request: "PipelineRequest", challenge: "HttpChallenge") -> None:
-        """authenticate according to challenge, add Authorization header to request"""
-
-        if self._need_new_token:
-            # azure-identity credentials require an AADv2 scope but the challenge may specify an AADv1 resource
-            scope = challenge.get_scope() or challenge.get_resource() + "/.default"
-            self._token = await self._credential.get_token(scope)
-
-        # ignore mypy's warning because although self._token is Optional, get_token raises when it fails to get a token
-        request.http_request.headers["Authorization"] = "Bearer {}".format(self._token.token)  # type: ignore
+        challenge = ChallengeCache.get_challenge_for_url(request.http_request.url)
+        if challenge:
+            # Note that if the vault has moved to a new tenant since our last request for it, this request will fail.
+            if self._need_new_token:
+                scope = _get_challenge_scope(challenge)
+                self._token = await self._credential.get_token(scope, tenant_id=challenge.tenant_id)
+
+            # ignore mypy's warning -- although self._token is Optional, get_token raises when it fails to get a token
+            request.http_request.headers["Authorization"] = "Bearer {}".format(self._token.token)  # type: ignore
+            return
+
+        # else: discover authentication information by eliciting a challenge from Key Vault. Remove any request data,
+        # saving it for later. Key Vault will reject the request as unauthorized and respond with a challenge.
+        # on_challenge will parse that challenge, reattach any body removed here, authorize the request, and tell
+        # super to send it again.
+        if request.http_request.body:
+            request.context["key_vault_request_data"] = request.http_request.body
+            request.http_request.set_json_body(None)
+            request.http_request.headers["Content-Length"] = "0"
+
+
+    async def on_challenge(self, request: "PipelineRequest", response: "PipelineResponse") -> bool:
+        try:
+            challenge = _update_challenge(request, response)
+            scope = _get_challenge_scope(challenge)
+        except ValueError:
+            return False
+
+        body = request.context.pop("key_vault_request_data", None)
+        request.http_request.set_text_body(body)  # no-op when text is None
+        await self.authorize_request(request, scope, tenant_id=challenge.tenant_id)
+
+        return True
+
+    @property
+    def _need_new_token(self) -> bool:
+        # pylint:disable=invalid-overridden-method
+        return not self._token or self._token.expires_on - time.time() < 300

sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/_internal/challenge_auth_policy.py

@@ -14,13 +14,11 @@
 protocol again.
 """
 
-import copy
 import time
 
 from azure.core.exceptions import ServiceRequestError
-from azure.core.pipeline import PipelineContext, PipelineRequest
-from azure.core.pipeline.policies import HTTPPolicy
-from azure.core.pipeline.transport import HttpRequest
+from azure.core.pipeline import PipelineRequest
+from azure.core.pipeline.policies import BearerTokenCredentialPolicy
 
 from .http_challenge import HttpChallenge
 from . import http_challenge_cache as ChallengeCache
@@ -44,20 +42,10 @@ def _enforce_tls(request):
         )
 
 
-def _get_challenge_request(request):
-    # type: (PipelineRequest) -> PipelineRequest
-
-    # The challenge request is intended to provoke an authentication challenge from Key Vault, to learn how the
-    # service request should be authenticated. It should be identical to the service request but with no body.
-    challenge_request = HttpRequest(
-        request.http_request.method, request.http_request.url, headers=request.http_request.headers
-    )
-    challenge_request.headers["Content-Length"] = "0"
-
-    options = copy.deepcopy(request.context.options)
-    context = PipelineContext(request.context.transport, **options)
-
-    return PipelineRequest(http_request=challenge_request, context=context)
+def _get_challenge_scope(challenge):
+    # type: (HttpChallenge) -> str
+    # azure-identity credentials require an AADv2 scope but the challenge may specify an AADv1 resource
+    return challenge.get_scope() or challenge.get_resource() + "/.default"
 
 
 def _update_challenge(request, challenger):
@@ -73,68 +61,53 @@ def _update_challenge(request, challenger):
     return challenge
 
 
-class ChallengeAuthPolicyBase(object):
-    """Sans I/O base for challenge authentication policies"""
+class ChallengeAuthPolicy(BearerTokenCredentialPolicy):
+    """policy for handling HTTP authentication challenges"""
 
-    def __init__(self, **kwargs):
+    def __init__(self, credential, *scopes, **kwargs):
+        # type: (TokenCredential, *str, **Any) -> None
+        super(ChallengeAuthPolicy, self).__init__(credential, *scopes, **kwargs)
+        self._credential = credential
         self._token = None  # type: Optional[AccessToken]
-        super(ChallengeAuthPolicyBase, self).__init__(**kwargs)
+
+    def on_request(self, request):
+        # type: (PipelineRequest) -> None
+        _enforce_tls(request)
+        challenge = ChallengeCache.get_challenge_for_url(request.http_request.url)
+        if challenge:
+            # Note that if the vault has moved to a new tenant since our last request for it, this request will fail.
+            if self._need_new_token:
+                scope = _get_challenge_scope(challenge)
+                self._token = self._credential.get_token(scope, tenant_id=challenge.tenant_id)
+
+            # ignore mypy's warning -- although self._token is Optional, get_token raises when it fails to get a token
+            request.http_request.headers["Authorization"] = "Bearer {}".format(self._token.token)  # type: ignore
+            return
+
+        # else: discover authentication information by eliciting a challenge from Key Vault. Remove any request data,
+        # saving it for later. Key Vault will reject the request as unauthorized and respond with a challenge.
+        # on_challenge will parse that challenge, reattach any body removed here, authorize the request, and tell
+        # super to send it again.
+        if request.http_request.body:
+            request.context["key_vault_request_data"] = request.http_request.body
+            request.http_request.set_json_body(None)
+            request.http_request.headers["Content-Length"] = "0"
+
+    def on_challenge(self, request, response):
+        # type: (PipelineRequest, PipelineResponse) -> bool
+        try:
+            challenge = _update_challenge(request, response)
+            scope = _get_challenge_scope(challenge)
+        except ValueError:
+            return False
+
+        body = request.context.pop("key_vault_request_data", None)
+        request.http_request.set_text_body(body)  # no-op when text is None
+        self.authorize_request(request, scope, tenant_id=challenge.tenant_id)
+
+        return True
 
     @property
     def _need_new_token(self):
         # type: () -> bool
         return not self._token or self._token.expires_on - time.time() < 300
-
-
-class ChallengeAuthPolicy(ChallengeAuthPolicyBase, HTTPPolicy):
-    """policy for handling HTTP authentication challenges"""
-
-    def __init__(self, credential, **kwargs):
-        # type: (TokenCredential, **Any) -> None
-        self._credential = credential
-        super(ChallengeAuthPolicy, self).__init__(**kwargs)
-
-    def send(self, request):
-        # type: (PipelineRequest) -> PipelineResponse
-        _enforce_tls(request)
-
-        challenge = ChallengeCache.get_challenge_for_url(request.http_request.url)
-        if not challenge:
-            challenge_request = _get_challenge_request(request)
-            challenger = self.next.send(challenge_request)
-            try:
-                challenge = _update_challenge(request, challenger)
-            except ValueError:
-                # didn't receive the expected challenge -> nothing more this policy can do
-                return challenger
-
-        self._handle_challenge(request, challenge)
-        response = self.next.send(request)
-
-        if response.http_response.status_code == 401:
-            # any cached token must be invalid
-            self._token = None
-
-            # cached challenge could be outdated; maybe this response has a new one?
-            try:
-                challenge = _update_challenge(request, response)
-            except ValueError:
-                # 401 with no legible challenge -> nothing more this policy can do
-                return response
-
-            self._handle_challenge(request, challenge)
-            response = self.next.send(request)
-
-        return response
-
-    def _handle_challenge(self, request, challenge):
-        # type: (PipelineRequest, HttpChallenge) -> None
-        """authenticate according to challenge, add Authorization header to request"""
-
-        if self._need_new_token:
-            # azure-identity credentials require an AADv2 scope but the challenge may specify an AADv1 resource
-            scope = challenge.get_scope() or challenge.get_resource() + "/.default"
-            self._token = self._credential.get_token(scope)
-
-        # ignore mypy's warning because although self._token is Optional, get_token raises when it fails to get a token
-        request.http_request.headers["Authorization"] = "Bearer {}".format(self._token.token)  # type: ignore

sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/_internal/http_challenge.py

@@ -40,6 +40,11 @@ def __init__(self, request_uri, challenge, response_headers=None):
         if "authorization" not in self._parameters and "authorization_uri" not in self._parameters:
             raise ValueError("Invalid challenge parameters")
 
+        authorization_uri = self.get_authorization_server()
+        # the authoritzation server URI should look something like https://login.windows.net/tenant-id
+        uri_path = parse.urlparse(authorization_uri).path.lstrip("/")
+        self.tenant_id = uri_path.split("/")[0] or None
+
         # if the response headers were supplied
         if response_headers:
             # get the message signing key and message key encryption key from the headers

sdk/keyvault/azure-keyvault-administration/setup.py

@@ -81,7 +81,7 @@
             "azure.keyvault",
         ]
     ),
-    install_requires=["azure-common~=1.1", "azure-core<2.0.0,>=1.11.0", "msrest>=0.6.21", "six>=1.11.0"],
+    install_requires=["azure-common~=1.1", "azure-core<2.0.0,>=1.15.0", "msrest>=0.6.21", "six>=1.11.0"],
     extras_require={
         ":python_version<'3.0'": ["azure-keyvault-nspkg"],
         ":python_version<'3.4'": ["enum34>=1.0.4"],

sdk/keyvault/azure-keyvault-certificates/CHANGELOG.md

@@ -3,12 +3,16 @@
 ## 4.4.0b2 (Unreleased)
 
 ### Features Added
+- Added support for multi-tenant authentication against Managed HSM when using
+  `azure-identity` 1.7.1 or newer
+  ([#20698](https://github.com/Azure/azure-sdk-for-python/issues/20698))
 
 ### Breaking Changes
 
 ### Bugs Fixed
 
 ### Other Changes
+- Updated minimum `azure-core` version to 1.15.0
 
 ## 4.4.0b1 (2021-09-09)
 

sdk/keyvault/azure-keyvault-certificates/azure/keyvault/certificates/_shared/__init__.py

@@ -9,7 +9,7 @@
     import urlparse as parse  # type: ignore
 
 from typing import TYPE_CHECKING
-from .challenge_auth_policy import ChallengeAuthPolicy, ChallengeAuthPolicyBase
+from .challenge_auth_policy import ChallengeAuthPolicy
 from .client_base import KeyVaultClientBase
 from .http_challenge import HttpChallenge
 from . import http_challenge_cache as HttpChallengeCache
@@ -21,7 +21,6 @@
 
 __all__ = [
     "ChallengeAuthPolicy",
-    "ChallengeAuthPolicyBase",
     "HttpChallenge",
     "HttpChallengeCache",
     "KeyVaultClientBase",