def test_bearer_challenge(self):
mock_bearer_challenge = ' Bearer authorization="https://login.windows.net/mock-id", resource="https://vault.azure.net"'
self.assertTrue(
HttpBearerChallenge.is_bearer_challenge(mock_bearer_challenge))
self.assertFalse(
HttpBearerChallenge.is_bearer_challenge('Some other string'))
with self.assertRaises(ValueError):
HttpBearerChallenge('https://test.uri.com', 'Non bearer string')
with self.assertRaises(ValueError):
HttpBearerChallenge('ftp://test.ftp.com', mock_bearer_challenge)
challenge = HttpBearerChallenge('https://test.uri.com',
mock_bearer_challenge)
self.assertEqual(challenge.get_authorization_server(),
'https://login.windows.net/mock-id')
self.assertEqual(challenge.get_resource(), 'https://vault.azure.net')
self.assertEqual(challenge.get_value('resource'),
'https://vault.azure.net')
self.assertEqual(challenge.get_scope(), '')
mock_bearer_challenge = ' Bearer authorization_uri="https://login.windows.net/mock-id", resource="https://vault.azure.net"'
challenge = HttpBearerChallenge('https://test.uri.com',
mock_bearer_challenge)
self.assertEqual(challenge.get_authorization_server(),
'https://login.windows.net/mock-id')
def mock_key_vault_auth_base(self, request):
challenge = HttpBearerChallenge(
request.url,
'Bearer authorization=fake-url,resource=https://vault.azure.net'
)
self.set_authorization_header(request, challenge)
return request
def handle_401(self, response, **kwargs):
"""
Takes the response authenticates and resends if neccissary
:return: The final response to the authenticated request
:rtype: requests.Response
"""
# If response is not 401 do not auth and return response
if not response.status_code == 401:
self._thread_local.auth_attempted = False
return response
auth_header = response.headers.get('www-authenticate', '')
# if the response auth header is not a bearer challenge do not auth and return response
if not HttpBearerChallenge.is_bearer_challenge(auth_header):
self._thread_local.auth_attempted = False
return response
# If we've already attempted to auth for this request once, do not auth and return response
if self._thread_local.auth_attempted:
self._thread_local.auth_attempted = False
return response
# Otherwise authenticate and retry the request
self._thread_local.auth_attempted = True
if self._thread_local.pos is not None:
# Rewind the file position indicator of the body to where
# it was to resend the request.
response.request.body.seek(self._thread_local.pos)
# add the challenge to the cache
challenge = HttpBearerChallenge(response.request.url, auth_header)
ChallengeCache.set_challenge_for_url(response.request.url, challenge)
# Consume content and release the original connection
# to allow our new request to reuse the same one.
response.content
response.close()
# copy the request to resend
prep = response.request.copy()
extract_cookies_to_jar(prep._cookies, response.request, response.raw)
prep.prepare_cookies(prep._cookies)
# setup the auth header on the copied request
self.set_authorization_header(prep, challenge)
# resend the request with proper authentication
_response = response.connection.send(prep, **kwargs)
_response.history.append(response)
_response.request = prep
return _response
def __call__(self, request):
# attempt to pre-fetch challenge if cached
if self._callback:
challenge = ChallengeCache.get_challenge_for_url(request.url)
if challenge:
# if challenge cached, retrieve token and update the request
self.set_authorization_header(request, challenge)
else:
# send the request to retrieve the challenge, then request the token and update
# the request
# TODO: wire up commons flag for things like Fiddler, logging, etc.
response = requests.Session().send(request)
if response.status_code == 401:
auth_header = response.headers['WWW-Authenticate']
if HttpBearerChallenge.is_bearer_challenge(auth_header):
challenge = HttpBearerChallenge(
response.request.url, auth_header)
ChallengeCache.set_challenge_for_url(
response.request.url, challenge)
self.set_authorization_header(request, challenge)
return request