Add challenge authentication to azure-keyvault-keys

sdk/keyvault/azure-keyvault-keys/azure/keyvault/keys/_internal.py

@@ -6,7 +6,6 @@
 from typing import TYPE_CHECKING
 from azure.core import Configuration
 from azure.core.pipeline import Pipeline
-from azure.core.pipeline.policies import BearerTokenCredentialPolicy
 from azure.core.pipeline.transport import RequestsTransport
 from ._generated import KeyVaultClient
 
@@ -21,6 +20,7 @@
 except ImportError:
     import urlparse as parse  # pylint: disable=import-error
 
+from ._shared.challenge_auth_policy import ChallengeAuthPolicy
 
 _VaultId = namedtuple("VaultId", ["vault_url", "collection", "name", "version"])
 
@@ -65,7 +65,7 @@ def create_config(credential, api_version=None, **kwargs):
         if api_version is None:
             api_version = KeyVaultClient.DEFAULT_API_VERSION
         config = KeyVaultClient.get_configuration_class(api_version, aio=False)(credential, **kwargs)
-        config.authentication_policy = BearerTokenCredentialPolicy(credential, KEY_VAULT_SCOPE)
+        config.authentication_policy = ChallengeAuthPolicy(credential)
         return config
 
     def __init__(self, vault_url, credential, config=None, transport=None, api_version=None, **kwargs):

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

@@ -0,0 +1,16 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+from .challenge_auth_policy import ChallengeAuthPolicy, ChallengeAuthPolicyBase
+from .http_challenge import HttpChallenge
+from . import http_challenge_cache as HttpChallengeCache
+
+__all__ = ["ChallengeAuthPolicy", "ChallengeAuthPolicyBase", "HttpChallenge", "HttpChallengeCache"]
+
+try:
+    from .async_challenge_auth_policy import AsyncChallengeAuthPolicy
+
+    __all__.append("AsyncChallengeAuthPolicy")
+except (SyntaxError, ImportError):
+    pass

sdk/keyvault/azure-keyvault-keys/azure/keyvault/keys/_shared/async_challenge_auth_policy.py

@@ -0,0 +1,57 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+from azure.core.pipeline import PipelineRequest
+from azure.core.pipeline.policies import AsyncHTTPPolicy
+from azure.core.pipeline.transport import HttpRequest, HttpResponse
+
+from . import ChallengeAuthPolicyBase, HttpChallenge, HttpChallengeCache
+
+
+class AsyncChallengeAuthPolicy(ChallengeAuthPolicyBase, AsyncHTTPPolicy):
+    """policy for handling HTTP authentication challenges"""
+
+    async def send(self, request: PipelineRequest) -> HttpResponse:
+        challenge = HttpChallengeCache.get_challenge_for_url(request.http_request.url)
+        if not challenge:
+            # provoke a challenge with an unauthorized, bodiless request
+            no_body = HttpRequest(
+                request.http_request.method, request.http_request.url, headers=request.http_request.headers
+            )
+            if request.http_request.body:
+                # no_body was created with request's headers -> if request has a body, no_body's content-length is wrong
+                no_body.headers["Content-Length"] = "0"
+
+            challenger = await self.next.send(PipelineRequest(http_request=no_body, context=request.context))
+            try:
+                challenge = self._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:
+            # cached challenge could be outdated; maybe this response has a new one?
+            try:
+                challenge = self._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"""
+
+        scope = challenge.get_resource()
+        if not scope.endswith("/.default"):
+            scope += "/.default"
+
+        access_token = await self._credential.get_token(scope)
+        self._update_headers(request.http_request.headers, access_token.token)

sdk/keyvault/azure-keyvault-keys/azure/keyvault/keys/_shared/challenge_auth_policy.py

@@ -0,0 +1,91 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+try:
+    from typing import TYPE_CHECKING
+except ImportError:
+    TYPE_CHECKING = False
+
+if TYPE_CHECKING:
+    # pylint:disable=unused-import
+    from azure.core.pipeline.transport import HttpResponse
+
+from azure.core.pipeline import PipelineRequest
+from azure.core.pipeline.policies import HTTPPolicy
+from azure.core.pipeline.policies.authentication import _BearerTokenCredentialPolicyBase
+from azure.core.pipeline.transport import HttpRequest
+
+from .http_challenge import HttpChallenge
+from . import http_challenge_cache as ChallengeCache
+
+
+class ChallengeAuthPolicyBase(_BearerTokenCredentialPolicyBase):
+    """Sans I/O base for challenge authentication policies"""
+
+    def __init__(self, credential, **kwargs):
+        super(ChallengeAuthPolicyBase, self).__init__(credential, **kwargs)
+
+    @staticmethod
+    def _update_challenge(request, challenger):
+        # type: (HttpRequest, HttpResponse) -> HttpChallenge
+        """parse challenge from challenger, cache it, return it"""
+
+        challenge = HttpChallenge(
+            request.http_request.url,
+            challenger.http_response.headers.get("WWW-Authenticate"),
+            response_headers=challenger.http_response.headers,
+        )
+        ChallengeCache.set_challenge_for_url(request.http_request.url, challenge)
+        return challenge
+
+
+class ChallengeAuthPolicy(ChallengeAuthPolicyBase, HTTPPolicy):
+    """policy for handling HTTP authentication challenges"""
+
+    def send(self, request):
+        # type: (PipelineRequest) -> HttpResponse
+
+        challenge = ChallengeCache.get_challenge_for_url(request.http_request.url)
+        if not challenge:
+            # provoke a challenge with an unauthorized, bodiless request
+            no_body = HttpRequest(
+                request.http_request.method, request.http_request.url, headers=request.http_request.headers
+            )
+            if request.http_request.body:
+                # no_body was created with request's headers -> if request has a body, no_body's content-length is wrong
+                no_body.headers["Content-Length"] = "0"
+
+            challenger = self.next.send(PipelineRequest(http_request=no_body, context=request.context))
+            try:
+                challenge = self._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:
+            # cached challenge could be outdated; maybe this response has a new one?
+            try:
+                challenge = self._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"""
+
+        scope = challenge.get_resource()
+        if not scope.endswith("/.default"):
+            scope += "/.default"
+
+        access_token = self._credential.get_token(scope)
+        self._update_headers(request.http_request.headers, access_token.token)

sdk/keyvault/azure-keyvault-keys/azure/keyvault/keys/_shared/http_challenge.py

@@ -0,0 +1,113 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+try:
+    import urllib.parse as parse
+except ImportError:
+    import urlparse as parse  # type: ignore
+
+
+class HttpChallenge(object):
+    def __init__(self, request_uri, challenge, response_headers=None):
+        """ Parses an HTTP WWW-Authentication Bearer challenge from a server. """
+        self.source_authority = self._validate_request_uri(request_uri)
+        self.source_uri = request_uri
+        self._parameters = {}
+
+        # get the scheme of the challenge and remove from the challenge string
+        trimmed_challenge = self._validate_challenge(challenge)
+        split_challenge = trimmed_challenge.split(" ", 1)
+        self.scheme = split_challenge[0]
+        trimmed_challenge = split_challenge[1]
+
+        # split trimmed challenge into comma-separated name=value pairs. Values are expected
+        # to be surrounded by quotes which are stripped here.
+        for item in trimmed_challenge.split(","):
+            # process name=value pairs
+            comps = item.split("=")
+            if len(comps) == 2:
+                key = comps[0].strip(' "')
+                value = comps[1].strip(' "')
+                if key:
+                    self._parameters[key] = value
+
+        # minimum set of parameters
+        if not self._parameters:
+            raise ValueError("Invalid challenge parameters")
+
+        # must specify authorization or authorization_uri
+        if "authorization" not in self._parameters and "authorization_uri" not in self._parameters:
+            raise ValueError("Invalid challenge parameters")
+
+        # if the response headers were supplied
+        if response_headers:
+            # get the message signing key and message key encryption key from the headers
+            self.server_signature_key = response_headers.get("x-ms-message-signing-key", None)
+            self.server_encryption_key = response_headers.get("x-ms-message-encryption-key", None)
+
+    def is_bearer_challenge(self):
+        """ Tests whether the HttpChallenge a Bearer challenge.
+        rtype: bool """
+        if not self.scheme:
+            return False
+
+        return self.scheme.lower() == "bearer"
+
+    def is_pop_challenge(self):
+        """ Tests whether the HttpChallenge is a proof of possession challenge.
+        rtype: bool """
+        if not self.scheme:
+            return False
+
+        return self.scheme.lower() == "pop"
+
+    def get_value(self, key):
+        return self._parameters.get(key)
+
+    def get_authorization_server(self):
+        """ Returns the URI for the authorization server if present, otherwise empty string. """
+        value = ""
+        for key in ["authorization_uri", "authorization"]:
+            value = self.get_value(key) or ""
+            if value:
+                break
+        return value
+
+    def get_resource(self):
+        """ Returns the resource if present, otherwise empty string. """
+        return self.get_value("resource") or ""
+
+    def get_scope(self):
+        """ Returns the scope if present, otherwise empty string. """
+        return self.get_value("scope") or ""
+
+    def supports_pop(self):
+        """ Returns True if challenge supports pop token auth else False """
+        return self._parameters.get("supportspop", "").lower() == "true"
+
+    def supports_message_protection(self):
+        """ Returns True if challenge vault supports message protection """
+        return self.supports_pop() and self.server_encryption_key and self.server_signature_key
+
+    def _validate_challenge(self, challenge):
+        """ Verifies that the challenge is a valid auth challenge and returns the key=value pairs. """
+        if not challenge:
+            raise ValueError("Challenge cannot be empty")
+
+        return challenge.strip()
+
+    # pylint: disable=no-self-use
+    def _validate_request_uri(self, uri):
+        """ Extracts the host authority from the given URI. """
+        if not uri:
+            raise ValueError("request_uri cannot be empty")
+
+        uri = parse.urlparse(uri)
+        if not uri.netloc:
+            raise ValueError("request_uri must be an absolute URI")
+
+        if uri.scheme.lower() not in ["http", "https"]:
+            raise ValueError("request_uri must be HTTP or HTTPS")
+
+        return uri.netloc

sdk/keyvault/azure-keyvault-keys/azure/keyvault/keys/_shared/http_challenge_cache.py

@@ -0,0 +1,89 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+import threading
+
+try:
+    import urllib.parse as parse
+except ImportError:
+    import urlparse as parse  # type: ignore
+
+try:
+    from typing import TYPE_CHECKING
+except ImportError:
+    TYPE_CHECKING = False
+
+if TYPE_CHECKING:
+    # pylint: disable=unused-import
+    from typing import Dict
+    from .http_challenge import HttpChallenge
+
+
+_cache = {}  # type: Dict[str, HttpChallenge]
+_lock = threading.Lock()
+
+
+def get_challenge_for_url(url):
+    """ Gets the challenge for the cached URL.
+    :param url: the URL the challenge is cached for.
+    :rtype: HttpBearerChallenge """
+
+    if not url:
+        raise ValueError("URL cannot be None")
+
+    key = _get_cache_key(url)
+
+    with _lock:
+        return _cache.get(key)
+
+
+def _get_cache_key(url):
+    """Use the URL's netloc as cache key except when the URL specifies the default port for its scheme. In that case
+    use the netloc without the port. That is to say, https://foo.bar and https://foo.bar:443 are considered equivalent.
+
+    This equivalency prevents an unnecessary challenge when using Key Vault's paging API. The Key Vault client doesn't
+    specify ports, but Key Vault's next page links do, so a redundant challenge would otherwise be executed when the
+    client requests the next page."""
+
+    parsed = parse.urlparse(url)
+    if parsed.scheme == "https" and parsed.port == 443:
+        return parsed.netloc[:-4]
+    return parsed.netloc
+
+
+def remove_challenge_for_url(url):
+    """ Removes the cached challenge for the specified URL.
+    :param url: the URL for which to remove the cached challenge """
+    if not url:
+        raise ValueError("URL cannot be empty")
+
+    url = parse.urlparse(url)
+
+    with _lock:
+        del _cache[url.netloc]
+
+
+def set_challenge_for_url(url, challenge):
+    """ Caches the challenge for the specified URL.
+    :param url: the URL for which to cache the challenge
+    :param challenge: the challenge to cache """
+    if not url:
+        raise ValueError("URL cannot be empty")
+
+    if not challenge:
+        raise ValueError("Challenge cannot be empty")
+
+    src_url = parse.urlparse(url)
+    if src_url.netloc != challenge.source_authority:
+        raise ValueError("Source URL and Challenge URL do not match")
+
+    with _lock:
+        _cache[src_url.netloc] = challenge
+
+
+def clear():
+    """ Clears the cache. """
+
+    with _lock:
+        _cache.clear()