1212 from urlparse import urlparse # type: ignore
1313
1414from azure .core .exceptions import ClientAuthenticationError
15- from azure .core .pipeline .policies import SansIOHTTPPolicy
15+ from azure .core .pipeline .policies import BearerTokenCredentialPolicy , SansIOHTTPPolicy
1616
1717try :
1818 from azure .core .pipeline .transport import AsyncHttpTransport
3333)
3434
3535if TYPE_CHECKING :
36- from azure .core .pipeline import PipelineRequest # pylint: disable=ungrouped-imports
36+ from typing import Any
37+ from azure .core .credentials import TokenCredential
38+ from azure .core .pipeline import PipelineResponse , PipelineRequest # pylint: disable=ungrouped-imports
3739
3840
3941class AzureSigningError (ClientAuthenticationError ):
@@ -44,6 +46,47 @@ class AzureSigningError(ClientAuthenticationError):
4446 """
4547
4648
49+ class _HttpChallenge (object ): # pylint:disable=too-few-public-methods
50+ """Represents a parsed HTTP WWW-Authentication Bearer challenge from a server."""
51+
52+ def __init__ (self , challenge ):
53+ if not challenge :
54+ raise ValueError ("Challenge cannot be empty" )
55+
56+ self ._parameters = {}
57+
58+ # Split the scheme ("Bearer") from the challenge parameters
59+ trimmed_challenge = challenge .strip ()
60+ split_challenge = trimmed_challenge .split (" " , 1 )
61+ trimmed_challenge = split_challenge [1 ]
62+
63+ # Split trimmed challenge into name=value pairs; these pairs are expected to be split by either commas or spaces
64+ # Values may be surrounded by quotes, which are stripped here
65+ separator = "," if "," in trimmed_challenge else " "
66+ for item in trimmed_challenge .split (separator ):
67+ # Process 'name=value' pairs
68+ comps = item .split ("=" )
69+ if len (comps ) == 2 :
70+ key = comps [0 ].strip (' "' )
71+ value = comps [1 ].strip (' "' )
72+ if key :
73+ self ._parameters [key ] = value
74+
75+ # Challenge must specify authorization or authorization_uri
76+ if not self ._parameters or (
77+ "authorization" not in self ._parameters and "authorization_uri" not in self ._parameters
78+ ):
79+ raise ValueError ("Invalid challenge parameters. `authorization` or `authorization_uri` must be present." )
80+
81+ authorization_uri = self ._parameters .get ("authorization" ) or self ._parameters .get ("authorization_uri" ) or ""
82+ # the authorization server URI should look something like https://login.windows.net/tenant-id[/oauth2/authorize]
83+ uri_path = urlparse (authorization_uri ).path .lstrip ("/" )
84+ self .tenant_id = uri_path .split ("/" )[0 ] or None
85+
86+ self .scope = self ._parameters .get ("scope" ) or ""
87+ self .resource = self ._parameters .get ("resource" ) or self ._parameters .get ("resource_id" ) or ""
88+
89+
4790# pylint: disable=no-self-use
4891class SharedKeyCredentialPolicy (SansIOHTTPPolicy ):
4992 def __init__ (self , credential , is_emulated = False ):
@@ -128,3 +171,60 @@ def _get_canonicalized_resource_query(self, request):
128171 if name == "comp" :
129172 return "?comp=" + value
130173 return ""
174+
175+
176+ class BearerTokenChallengePolicy (BearerTokenCredentialPolicy ):
177+ """Adds a bearer token Authorization header to requests, for the tenant provided in authentication challenges.
178+
179+ See https://docs.microsoft.com/azure/active-directory/develop/claims-challenge for documentation on AAD
180+ authentication challenges.
181+
182+ :param credential: The credential.
183+ :type credential: ~azure.core.TokenCredential
184+ :param str scopes: Lets you specify the type of access needed.
185+ :keyword bool discover_tenant: Determines if tenant discovery should be enabled. Defaults to True.
186+ :keyword bool discover_scopes: Determines if scopes from authentication challenges should be provided to token
187+ requests, instead of the scopes given to the policy's constructor, if any are present. Defaults to True.
188+ :raises: :class:`~azure.core.exceptions.ServiceRequestError`
189+ """
190+
191+ def __init__ (
192+ self ,
193+ credential : "TokenCredential" ,
194+ * scopes : str ,
195+ discover_tenant : bool = True ,
196+ discover_scopes : bool = True ,
197+ ** kwargs : "Any"
198+ ) -> None :
199+ self ._discover_tenant = discover_tenant
200+ self ._discover_scopes = discover_scopes
201+ super ().__init__ (credential , * scopes , ** kwargs )
202+
203+ def on_challenge (self , request : "PipelineRequest" , response : "PipelineResponse" ) -> bool :
204+ """Authorize request according to an authentication challenge
205+
206+ This method is called when the resource provider responds 401 with a WWW-Authenticate header.
207+
208+ :param ~azure.core.pipeline.PipelineRequest request: the request which elicited an authentication challenge
209+ :param ~azure.core.pipeline.PipelineResponse response: the resource provider's response
210+ :returns: a bool indicating whether the policy should send the request
211+ """
212+ if not self ._discover_tenant and not self ._discover_scopes :
213+ # We can't discover the tenant or use a different scope; the request will fail because it hasn't changed
214+ return False
215+
216+ try :
217+ challenge = _HttpChallenge (response .http_response .headers .get ("WWW-Authenticate" ))
218+ # azure-identity credentials require an AADv2 scope but the challenge may specify an AADv1 resource
219+ # if no scopes are included in the challenge, challenge.scope and challenge.resource will both be ''
220+ scope = challenge .scope or challenge .resource + "/.default" if self ._discover_scopes else self ._scopes
221+ if scope == "/.default" :
222+ scope = self ._scopes
223+ except ValueError :
224+ return False
225+
226+ if self ._discover_tenant :
227+ self .authorize_request (request , scope , tenant_id = challenge .tenant_id )
228+ else :
229+ self .authorize_request (request , scope )
230+ return True
0 commit comments