[Core] Add authorization challenge parsing helper

sdk/core/azure-core/azure/core/pipeline/policies/__init__.py

@@ -25,7 +25,12 @@
 # --------------------------------------------------------------------------
 
 from ._base import HTTPPolicy, SansIOHTTPPolicy, RequestHistory
-from ._authentication import BearerTokenCredentialPolicy, AzureKeyCredentialPolicy, AzureSasCredentialPolicy
+from ._authentication import (
+    AuthorizationChallengeParser,
+    BearerTokenCredentialPolicy,
+    AzureKeyCredentialPolicy,
+    AzureSasCredentialPolicy,
+)
 from ._custom_hook import CustomHookPolicy
 from ._redirect import RedirectPolicy
 from ._retry import RetryPolicy, RetryMode
@@ -43,6 +48,7 @@
 __all__ = [
     'HTTPPolicy',
     'SansIOHTTPPolicy',
+    'AuthorizationChallengeParser',
     'BearerTokenCredentialPolicy',
     'AzureKeyCredentialPolicy',
     'AzureSasCredentialPolicy',

sdk/core/azure-core/azure/core/pipeline/policies/_authentication.py

@@ -21,6 +21,52 @@
     from azure.core.pipeline import PipelineRequest, PipelineResponse
 
 
+class AuthorizationChallengeParser(object):  # pylint:disable=too-few-public-methods
+    """A helper class for parsing Authorization challenge headers."""
+
+    @staticmethod
+    def get_challenge_parameter(response: "PipelineResponse", scheme: str, parameter: str) -> "Optional[str]":
+        """Parses the specified challenge parameter from the WWW-Authenticate header of the provided response.
+
+        :param response: The challenge response.
+        :type response: ~azure.core.pipeline.PipelineResponse
+        :param str scheme: The challenge scheme containing the challenge parameter. For example, "Bearer".
+        :param str parameter: The parameter key name containing the value to return. For example, "resource".
+
+        :returns: The string value of the challenge parameter if it's present, and None otherwise.
+        :rtype: str or None
+        :raises: ValueError if the challenge is improperly formatted.
+        """
+
+        challenge_header = response.http_response.headers.get("WWW-Authenticate")
+        if not challenge_header:
+            raise ValueError("No WWW-Authenticate header found in the response; challenge cannot be empty.")
+
+        # Split the scheme (e.g. "Bearer") from the challenge parameters
+        trimmed_challenge = challenge_header.strip()
+        split_challenge = trimmed_challenge.split(" ", 1)
+        challenge_scheme = split_challenge[0]
+        # Return early if the challenge scheme doesn't match the queried scheme
+        if challenge_scheme.lower() != scheme.lower():
+            return None
+
+        # Split trimmed challenge into name=value pairs; these pairs are expected to be split by either commas or spaces
+        # Values may be surrounded by quotes, which are stripped here
+        trimmed_challenge = split_challenge[1]
+        parameters = {}
+        separator = "," if "," in trimmed_challenge else " "
+        for item in trimmed_challenge.split(separator):
+            # Process 'name=value' pairs
+            comps = item.split("=")
+            if len(comps) == 2:
+                key = comps[0].strip(' "')
+                value = comps[1].strip(' "')
+                if key:
+                    parameters[key.lower()] = value
+
+        return parameters.get(parameter.lower())
+
+
 # pylint:disable=too-few-public-methods
 class _BearerTokenCredentialPolicyBase(object):
     """Base class for a Bearer Token Credential Policy.
@@ -180,6 +226,7 @@ class AzureKeyCredentialPolicy(SansIOHTTPPolicy):
     :param str name: The name of the key header used for the credential.
     :raises: ValueError or TypeError
     """
+
     def __init__(self, credential, name, **kwargs):  # pylint: disable=unused-argument
         # type: (AzureKeyCredential, str, **Any) -> None
         super(AzureKeyCredentialPolicy, self).__init__()
@@ -201,6 +248,7 @@ class AzureSasCredentialPolicy(SansIOHTTPPolicy):
     :type credential: ~azure.core.credentials.AzureSasCredential
     :raises: ValueError or TypeError
     """
+
     def __init__(self, credential, **kwargs):  # pylint: disable=unused-argument
         # type: (AzureSasCredential, **Any) -> None
         super(AzureSasCredentialPolicy, self).__init__()

sdk/core/azure-core/tests/test_authentication.py

@@ -10,6 +10,7 @@
 from azure.core.exceptions import ServiceRequestError
 from azure.core.pipeline import Pipeline
 from azure.core.pipeline.policies import (
+    AuthorizationChallengeParser,
     BearerTokenCredentialPolicy,
     SansIOHTTPPolicy,
     AzureKeyCredentialPolicy,
@@ -389,3 +390,60 @@ def test_azure_named_key_credential_raises():
 
     with pytest.raises(TypeError, match="Both name and key must be strings."):
         cred.update(1234, "newkey")
+
+def test_authorization_challenge_parser():
+    endpoint = f"https://authority.net/tenant-id/oauth2/authorize"
+    resource = "https://challenge.resource"
+    scope = f"{resource}/.default"
+
+    # this challenge separates the authorization server and resource with commas in the WWW-Authenticate header
+    challenge_with_commas = Mock(
+        status_code=401,
+        headers={"WWW-Authenticate": f'Bearer authorization="{endpoint}", resource="{resource}"'},
+    )
+    response_with_commas = Mock(http_response=challenge_with_commas)
+
+    challenge_authorization = AuthorizationChallengeParser.get_challenge_parameter(
+        response_with_commas, "Bearer", "authorization"
+    )
+    challenge_resource = AuthorizationChallengeParser.get_challenge_parameter(
+        response_with_commas, "Bearer", "resource"
+    )
+    challenge_scope = AuthorizationChallengeParser.get_challenge_parameter(response_with_commas, "Bearer", "scope")
+    assert challenge_authorization == endpoint
+    assert challenge_resource == resource
+    assert challenge_scope is None
+
+    # this challenge separates the authorization server and resource with only spaces in the WWW-Authenticate header
+    challenge_without_commas = Mock(
+        status_code=401,
+        headers={"WWW-Authenticate": f'Bearer authorization={endpoint} resource={resource}'},
+    )
+    response_without_commas = Mock(http_response=challenge_without_commas)
+
+    challenge_authorization = AuthorizationChallengeParser.get_challenge_parameter(
+        response_without_commas, "BEARER", "AUTHORIZATION"  # scheme and parameter should each be case-insensitive
+    )
+    challenge_resource = AuthorizationChallengeParser.get_challenge_parameter(
+        response_without_commas, "Bearer", "resource"
+    )
+    challenge_scope = AuthorizationChallengeParser.get_challenge_parameter(response_without_commas, "Bearer", "scope")
+    assert challenge_authorization == endpoint
+    assert challenge_resource == resource
+    assert challenge_scope is None
+
+    # this challenge gives an AADv2 scope, ending with "/.default", instead of an AADv1 resource
+    challenge_with_scope = Mock(
+        status_code=401,
+        headers={"WWW-Authenticate": f'Bearer authorization={endpoint} scope={scope}'},
+    )
+    response_with_scope = Mock(http_response=challenge_with_scope)
+
+    challenge_authorization = AuthorizationChallengeParser.get_challenge_parameter(
+        response_with_scope, "Bearer", "authorization"
+    )
+    challenge_resource = AuthorizationChallengeParser.get_challenge_parameter(response_with_scope, "Bearer", "resource")
+    challenge_scope = AuthorizationChallengeParser.get_challenge_parameter(response_with_scope, "Bearer", "scope")
+    assert challenge_authorization == endpoint
+    assert challenge_resource is None
+    assert challenge_scope == scope