89 lines
3.3 KiB
Diff
89 lines
3.3 KiB
Diff
diff --git a/framework/replay/download_utils.py b/framework/replay/download_utils.py
|
|
index 36322b000..5c3fe140d 100644
|
|
--- a/framework/replay/download_utils.py
|
|
+++ b/framework/replay/download_utils.py
|
|
@@ -27,20 +27,20 @@ import base64
|
|
import hashlib
|
|
import hmac
|
|
import xml.etree.ElementTree as ET
|
|
-
|
|
-from typing import Dict
|
|
from email.utils import formatdate
|
|
from os import path
|
|
from time import time
|
|
+from typing import Dict
|
|
+from urllib.parse import urlparse
|
|
+
|
|
import requests
|
|
from requests.adapters import HTTPAdapter, Retry
|
|
-from framework.replay.local_file_adapter import LocalFileAdapter
|
|
from requests.utils import requote_uri
|
|
|
|
from framework import core, exceptions
|
|
+from framework.replay.local_file_adapter import LocalFileAdapter
|
|
from framework.replay.options import OPTIONS
|
|
|
|
-
|
|
__all__ = ['ensure_file']
|
|
|
|
minio_credentials = None
|
|
@@ -90,7 +90,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)
|
|
|
|
date = formatdate(timeval=None, localtime=False, usegmt=True)
|
|
@@ -107,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: Dict[str, str], attempts: int = 2) -> None:
|
|
"""Downloads a URL content into a file
|
|
|
|
@@ -178,7 +189,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'))
|