diff --git a/fades/helpers.py b/fades/helpers.py index 6aa7802..a8bee74 100644 --- a/fades/helpers.py +++ b/fades/helpers.py @@ -301,7 +301,21 @@ def _download_raw(self, url=None): if url is None: url = self.url req = request.Request(url, headers=self.HEADERS_PLAIN) - return request.urlopen(req).read().decode("utf8") + resp = request.urlopen(req) + + # check if the response url is different than the original one; in this case we had + # redirected, and we need to pass the new url response through the proper + # pastebin-dependant adapter, so recursively go into another _ScriptDownloader + if resp.geturl() != url: + new_url = resp.geturl() + downloader = _ScriptDownloader(new_url) + logger.info( + "Download redirect detect, now downloading from %r using %r downloader", + new_url, downloader.name) + return downloader.get() + + # simple non-redirect response + return resp.read().decode("utf8") def _download_linkode(self): """Download content from Linkode pastebin.""" @@ -336,7 +350,7 @@ def download_remote_script(url): temp_fh = tempfile.NamedTemporaryFile('wt', encoding='utf8', suffix=".py", delete=False) downloader = _ScriptDownloader(url) logger.info( - "Downloading remote script from %r using (%r downloader) to %r", + "Downloading remote script from %r (using %r downloader) to %r", url, downloader.name, temp_fh.name) content = downloader.get() diff --git a/tests/test_helpers.py b/tests/test_helpers.py index e522d27..d4232d8 100644 --- a/tests/test_helpers.py +++ b/tests/test_helpers.py @@ -398,7 +398,8 @@ def test_external_public_function(self): # checks mock_downloader_class.assert_called_with(test_url) self.assertLoggedInfo( - "Downloading remote script", test_url, repr(filepath), 'mock downloader') + "Downloading remote script from {!r}".format(test_url), + repr(filepath), "(using 'mock downloader' downloader)") with open(filepath, "rt", encoding='utf8') as fh: self.assertEqual(fh.read(), test_content) @@ -428,6 +429,7 @@ def test_downloader_raw(self): with patch('http.client.HTTPResponse') as mock_http_response: mock_http_response.read.return_value = raw_service_response mock_urlopen.return_value = mock_http_response + mock_http_response.geturl.return_value = test_url content = downloader.get() @@ -474,6 +476,7 @@ def test_downloader_linkode(self): def test_downloader_pastebin(self): test_url = "http://pastebin.com/sZGwz7SL" + real_url = "https://pastebin.com/raw/sZGwz7SL" test_content = "test content of the remote script áéíóú" raw_service_response = test_content.encode("utf8") @@ -482,6 +485,7 @@ def test_downloader_pastebin(self): with patch('http.client.HTTPResponse') as mock_http_response: mock_http_response.read.return_value = raw_service_response mock_urlopen.return_value = mock_http_response + mock_http_response.geturl.return_value = real_url content = downloader.get() @@ -493,13 +497,13 @@ def test_downloader_pastebin(self): (call,) = mock_urlopen.mock_calls (called_request,) = call[1] self.assertIsInstance(called_request, Request) - self.assertEqual( - called_request.full_url, "https://pastebin.com/raw/sZGwz7SL") + self.assertEqual(called_request.full_url, real_url) self.assertEqual(called_request.headers, headers) self.assertEqual(content, test_content) def test_downloader_gist(self): test_url = "http://gist.github.com/facundobatista/6ff4f75760a9acc35e68bae8c1d7da1c" + real_url = "https://gist.github.com/facundobatista/6ff4f75760a9acc35e68bae8c1d7da1c/raw" test_content = "test content of the remote script áéíóú" raw_service_response = test_content.encode("utf8") @@ -508,6 +512,7 @@ def test_downloader_gist(self): with patch('http.client.HTTPResponse') as mock_http_response: mock_http_response.read.return_value = raw_service_response mock_urlopen.return_value = mock_http_response + mock_http_response.geturl.return_value = real_url content = downloader.get() @@ -519,8 +524,35 @@ def test_downloader_gist(self): (call,) = mock_urlopen.mock_calls (called_request,) = call[1] self.assertIsInstance(called_request, Request) - self.assertEqual( - called_request.full_url, - "https://gist.github.com/facundobatista/6ff4f75760a9acc35e68bae8c1d7da1c/raw") + self.assertEqual(called_request.full_url, real_url) self.assertEqual(called_request.headers, headers) self.assertEqual(content, test_content) + + def test_downloader_raw_with_redirection(self): + test_url = "http://bit.ly/will-redirect" + final_url = "http://real-service.com/" + raw_service_response = b"test content of the remote script" + downloader = helpers._ScriptDownloader(test_url) + response_contents = [ + b"whatever; we don't care as we are redirectect", + raw_service_response, + ] + with patch('urllib.request.urlopen') as mock_urlopen: + with patch('http.client.HTTPResponse') as mock_http_response: + mock_http_response.read.side_effect = lambda: response_contents.pop() + mock_http_response.geturl.return_value = final_url + mock_urlopen.return_value = mock_http_response + + content = downloader.get() + + # two calls, first to the service that will redirect us, second to the final one + call1, call2 = mock_urlopen.mock_calls + + (called_request,) = call1[1] + self.assertEqual(called_request.full_url, test_url) + + (called_request,) = call2[1] + self.assertEqual(called_request.full_url, final_url) + self.assertEqual(content, raw_service_response.decode("utf8")) + + self.assertLoggedInfo("Download redirect detect, now downloading from", final_url)