diff --git a/framework/replay/download_utils.py b/framework/replay/download_utils.py index 3119a24a2..4e776ca85 100644 --- a/framework/replay/download_utils.py +++ b/framework/replay/download_utils.py @@ -31,6 +31,7 @@ import xml.etree.ElementTree as ET from email.utils import formatdate from os import path from time import time +from urllib.parse import urlparse import requests from requests.adapters import HTTPAdapter, Retry from requests.utils import requote_uri @@ -88,7 +89,7 @@ def get_minio_credentials(url): minio_credentials['SessionToken']) -def get_authorization_headers(url, resource): +def get_minio_authorization_headers(url, resource): minio_key, minio_secret, minio_token = get_minio_credentials(url) content_type = 'application/octet-stream' @@ -106,6 +107,17 @@ def get_authorization_headers(url, resource): return headers +def get_jwt_authorization_headers(url, resource): + date = formatdate(timeval=None, localtime=False, usegmt=True) + jwt = OPTIONS.download['jwt'] + host = urlparse(url).netloc + + headers = {'Host': host, + 'Date': date, + 'Authorization': 'Bearer %s' % (jwt)} + return headers + + def download(url: str, file_path: str, headers, attempts=2) -> None: """Downloads a URL content into a file @@ -174,7 +186,9 @@ def ensure_file(file_path): assert OPTIONS.download['minio_bucket'] assert OPTIONS.download['role_session_name'] assert OPTIONS.download['jwt'] - headers = get_authorization_headers(url, file_path) + headers = get_minio_authorization_headers(url, file_path) + elif OPTIONS.download['jwt']: + headers = get_jwt_authorization_headers(url, file_path) else: headers = None diff --git a/unittests/framework/replay/test_download_utils.py b/unittests/framework/replay/test_download_utils.py index 1e78b26e7..749c5d835 100644 --- a/unittests/framework/replay/test_download_utils.py +++ b/unittests/framework/replay/test_download_utils.py @@ -195,3 +195,17 @@ class TestDownloadUtils(object): get_request = requests_mock.request_history[1] assert(get_request.method == 'GET') assert(requests_mock.request_history[1].headers['Authorization'].startswith('AWS Key')) + + def test_jwt_authorization(self, requests_mock): + """download_utils.ensure_file: Check we send the authentication headers to the server""" + # reset minio_host from previous tests + OPTIONS.download['minio_host'] = '' + OPTIONS.download['jwt'] = 'jwt' + + assert not self.trace_file.check() + download_utils.ensure_file(self.trace_path) + TestDownloadUtils.check_same_file(self.trace_file, "remote") + + get_request = requests_mock.request_history[0] + assert(get_request.method == 'GET') + assert(requests_mock.request_history[0].headers['Authorization'].startswith('Bearer'))