Support download referer headers and redirects

This commit is contained in:
2Shirt 2022-09-27 23:50:13 -07:00
parent 9f52daeec3
commit 3b89f1eabc
Signed by: 2Shirt
GPG key ID: 152FAC923B0E132C
2 changed files with 19 additions and 7 deletions

View file

@ -75,10 +75,10 @@ def delete_from_temp(item_path):
delete_item(TMP_DIR.joinpath(item_path), force=True, ignore_errors=True)
def download_to_temp(filename, source_url):
def download_to_temp(filename, source_url, referer=None):
"""Download file to temp dir, returns pathlib.Path."""
out_path = TMP_DIR.joinpath(filename)
download_file(out_path, source_url)
download_file(out_path, source_url, referer=referer)
return out_path
@ -209,7 +209,7 @@ def download_fastcopy():
def download_furmark():
"""Download FurMark."""
installer = download_to_temp('FurMark_Setup.exe', SOURCES['FurMark'])
installer = download_to_temp('FurMark_Setup.exe', SOURCES['FurMark'], referer=SOURCES['FurMark'])
out_path = BIN_DIR.joinpath('FurMark')
tmp_path = TMP_DIR.joinpath('FurMarkInstall')
run_program([installer, f'/DIR={tmp_path}', '/SILENT'])

View file

@ -11,7 +11,7 @@ import requests
from wk.cfg.main import ARCHIVE_PASSWORD
from wk.cfg.sources import DOWNLOAD_FREQUENCY, SOURCES
from wk.exe import popen_program, run_program
from wk.std import GenericError
from wk.std import GenericError, sleep
# STATIC VARIABLES
@ -30,7 +30,7 @@ CACHED_DIRS = {}
# Functions
def download_file(out_path, source_url, as_new=False, overwrite=False):
def download_file(out_path, source_url, as_new=False, overwrite=False, referer=None):
"""Download a file using requests, returns pathlib.Path."""
out_path = pathlib.Path(out_path).resolve()
name = out_path.name
@ -47,12 +47,24 @@ def download_file(out_path, source_url, as_new=False, overwrite=False):
# Create destination directory
out_path.parent.mkdir(parents=True, exist_ok=True)
# Update headers
headers = HEADERS.copy()
if referer:
headers['referer'] = referer
# Request download
with requests.Session() as session:
try:
response = session.get(source_url, headers=HEADERS, stream=True)
response = session.get(source_url, allow_redirects=True, headers=headers, stream=True)
except requests.RequestException as _err:
download_failed = _err
try:
sleep(1)
response = session.get(source_url, allow_redirects=True, headers=headers, stream=True)
except requests.RequestException as _err:
download_failed = _err
else:
if not response.ok:
download_failed = response
else:
if not response.ok:
download_failed = response