def _get_path(self, path):
"""
Retrieve a single path within the upstream registry, and return a 2-tuple of the headers and
the response body.
:param path: a full http path to retrieve that will be urljoin'd to the upstream registry
url.
:type path: basestring
:return: (headers, response body)
:rtype: tuple
"""
url = urlparse.urljoin(self.registry_url, path)
_logger.debug(_('Retrieving {0}'.format(url)))
request = DownloadRequest(url, StringIO())
if self.token:
token_util.add_auth_header(request, self.token)
report = self.downloader.download_one(request)
# If the download was unauthorized, attempt to get a token and try again
if report.state == report.DOWNLOAD_FAILED:
if report.error_report.get('response_code') == httplib.UNAUTHORIZED:
_logger.debug(_('Download unauthorized, attempting to retrieve a token.'))
self.token = token_util.request_token(self.downloader, request, report.headers)
token_util.add_auth_header(request, self.token)
report = self.downloader.download_one(request)
if report.state == report.DOWNLOAD_FAILED:
raise IOError(report.error_msg)
return report.headers, report.destination.getvalue()
def test_with_headers(self):
"""
Test that when the headers exists, the auth token is added to it.
"""
mock_req = mock.MagicMock()
mock_req.headers = {"mock": "header"}
token_util.add_auth_header(mock_req, "mock token")
self.assertDictEqual(mock_req.headers, {"Authorization": "Bearer mock token", "mock": "header"})
def test_no_headers(self):
"""
Test that when there are no existing headers, it is added.
"""
mock_req = mock.MagicMock()
mock_req.headers = None
token_util.add_auth_header(mock_req, "mock token")
self.assertDictEqual(mock_req.headers, {"Authorization": "Bearer mock token"})
def test_no_headers(self):
"""
Test that when there are no existing headers, it is added.
"""
mock_req = mock.MagicMock()
mock_req.headers = None
token_util.add_auth_header(mock_req, "mock token")
self.assertDictEqual(mock_req.headers,
{"Authorization": "Bearer mock token"})
def process_main(self, item=None):
"""
Overrides the parent method to get a new token and try again if response is a 401.
"""
# Allow the original request to be retrieved from the url.
for request in self.downloads:
self._requests_map[request.url] = request
for request in self.downloads:
if self.token:
token_util.add_auth_header(request, self.token)
self.downloader.download_one(request, events=True)
def test_with_headers(self):
"""
Test that when the headers exists, the auth token is added to it.
"""
mock_req = mock.MagicMock()
mock_req.headers = {"mock": "header"}
token_util.add_auth_header(mock_req, "mock token")
self.assertDictEqual(mock_req.headers, {
"Authorization": "Bearer mock token",
"mock": "header"
})
def download_failed(self, report):
"""
If the download fails due to a 401, attempt to retreive a token and try again.
:param report: download report
:type report: nectar.report.DownloadReport
"""
if report.error_report.get('response_code') == httplib.UNAUTHORIZED:
_logger.debug(_('Download unauthorized, attempting to retrieve a token.'))
request = self._requests_map[report.url]
token = token_util.request_token(self.downloader, request, report.headers)
token_util.add_auth_header(request, token)
_logger.debug("Trying download again with new bearer token.")
report = self.downloader.download_one(request, events=True)
if report.state == report.DOWNLOAD_FAILED:
super(TokenAuthDownloadStep, self).download_failed(report)
def _get_path(self, path):
"""
Retrieve a single path within the upstream registry, and return a 2-tuple of the headers and
the response body.
:param path: a full http path to retrieve that will be urljoin'd to the upstream registry
url.
:type path: basestring
:return: (headers, response body)
:rtype: tuple
"""
url = urlparse.urljoin(self.registry_url, path)
_logger.debug(_('Retrieving {0}'.format(url)))
request = DownloadRequest(url, StringIO())
if self.token:
token_util.add_auth_header(request, self.token)
report = self.downloader.download_one(request)
# If the download was unauthorized, attempt to get a token and try again
if report.state == report.DOWNLOAD_FAILED:
if report.error_report.get(
'response_code') == httplib.UNAUTHORIZED:
_logger.debug(
_('Download unauthorized, attempting to retrieve a token.')
)
self.token = token_util.request_token(self.downloader, request,
report.headers)
token_util.add_auth_header(request, self.token)
report = self.downloader.download_one(request)
if report.state == report.DOWNLOAD_FAILED:
raise IOError(report.error_msg)
return report.headers, report.destination.getvalue()