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 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()