Fix redirect to s3

This commit is contained in:
Harrison Wright
2019-06-22 13:00:42 -05:00
parent de0c3f46c6
commit 9b6400932d

View File

@@ -26,6 +26,8 @@ try:
from urllib.error import HTTPError, URLError from urllib.error import HTTPError, URLError
from urllib.request import urlopen from urllib.request import urlopen
from urllib.request import Request from urllib.request import Request
from urllib.request import HTTPRedirectHandler
from urllib.request import build_opener
except ImportError: except ImportError:
# python 2 # python 2
from urlparse import urlparse from urlparse import urlparse
@@ -34,6 +36,8 @@ except ImportError:
from urllib2 import HTTPError, URLError from urllib2 import HTTPError, URLError
from urllib2 import urlopen from urllib2 import urlopen
from urllib2 import Request from urllib2 import Request
from urllib2 import HTTPRedirectHandler
from urllib2 import build_opener
from github_backup import __version__ from github_backup import __version__
@@ -537,22 +541,33 @@ def _request_url_error(template, retry_timeout):
return False return False
class S3HTTPRedirectHandler(HTTPRedirectHandler):
"""
A subclassed redirect handler for downloading Github assets from S3.
urllib will add the Authorization header to the redirected request to S3, which will result in a 400,
so we should remove said header on redirect.
"""
def redirect_request(self, req, fp, code, msg, headers, newurl):
request = super(S3HTTPRedirectHandler, self).redirect_request(req, fp, code, msg, headers, newurl)
del request.headers['Authorization']
return request
def download_file(url, path, auth): def download_file(url, path, auth):
request = Request(url) request = Request(url)
request.add_header('Accept', 'application/octet-stream') request.add_header('Accept', 'application/octet-stream')
request.add_header('Authorization', 'Basic '.encode('ascii') + auth) request.add_header('Authorization', 'Basic '.encode('ascii') + auth)
data = urlopen(request) opener = build_opener(S3HTTPRedirectHandler)
with open(path, 'wb') as f: response = opener.open(request)
f.write(data.read())
# import requests chunk_size = 16 * 1024
# r = requests.get(url, stream=True, headers={ with open(path, 'wb') as f:
# 'Accept': 'application/octet-stream', while True:
# 'Authorization': 'Basic '.encode('ascii') + auth chunk = response.read(chunk_size)
# }) if not chunk:
# with open(path, 'wb') as f: break
# for chunk in r.iter_content(1024): f.write(chunk)
# f.write(chunk)
def get_authenticated_user(args): def get_authenticated_user(args):