framework/replay/download_utils.py

   1 # coding=utf-8
   2 #
   3 # Copyright © 2020, 2022 Collabora Ltd
   4 # Copyright © 2020 Valve Corporation.
   5 #
   6 # Permission is hereby granted, free of charge, to any person obtaining a
   7 # copy of this software and associated documentation files (the "Software"),
   8 # to deal in the Software without restriction, including without limitation
   9 # the rights to use, copy, modify, merge, publish, distribute, sublicense,
  10 # and/or sell copies of the Software, and to permit persons to whom the
  11 # Software is furnished to do so, subject to the following conditions:
  12 #
  13 # The above copyright notice and this permission notice shall be included
  14 # in all copies or substantial portions of the Software.
  15 #
  16 # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
  17 # OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  18 # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
  19 # THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
  20 # OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
  21 # ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
  22 # OTHER DEALINGS IN THE SOFTWARE.
  23 #
  24 # SPDX-License-Identifier: MIT
  25
  26 import base64
  27 import hashlib
  28 import hmac
  29 import xml.etree.ElementTree as ET
  30
  31 from typing import Dict
  32 from email.utils import formatdate
  33 from os import path
  34 from time import time
  35 import requests
  36 from requests.adapters import HTTPAdapter, Retry
  37 from framework.replay.local_file_adapter import LocalFileAdapter
  38 from requests.utils import requote_uri
  39
  40 from framework import core, exceptions
  41 from framework.replay.options import OPTIONS
  42
  43
  44 __all__ = ['ensure_file']
  45
  46 minio_credentials = None
  47
  48
  49 def sign_with_hmac(key, message):
  50     key = key.encode("UTF-8")
  51     message = message.encode("UTF-8")
  52
  53     signature = hmac.new(key, message, hashlib.sha1).digest()
  54
  55     return base64.encodebytes(signature).strip().decode()
  56
  57
  58 def get_minio_credentials(url):
  59     global minio_credentials
  60
  61     if minio_credentials is not None:
  62         return (minio_credentials['AccessKeyId'],
  63                 minio_credentials['SecretAccessKey'],
  64                 minio_credentials['SessionToken'])
  65
  66     minio_credentials = {}
  67
  68     params = {'Action': 'AssumeRoleWithWebIdentity',
  69               'Version': '2011-06-15',
  70               'RoleArn': 'arn:aws:iam::123456789012:role/FederatedWebIdentityRole',
  71               'RoleSessionName': OPTIONS.download['role_session_name'],
  72               'DurationSeconds': 3600,
  73               'WebIdentityToken': OPTIONS.download['jwt']}
  74     r = requests.post('https://%s' % OPTIONS.download['minio_host'], params=params)
  75     if r.status_code >= 400:
  76         print(r.text)
  77     r.raise_for_status()
  78
  79     root = ET.fromstring(r.text)
  80     for attr in root.iter():
  81         if attr.tag == '{https://sts.amazonaws.com/doc/2011-06-15/}AccessKeyId':
  82             minio_credentials['AccessKeyId'] = attr.text
  83         elif attr.tag == '{https://sts.amazonaws.com/doc/2011-06-15/}SecretAccessKey':
  84             minio_credentials['SecretAccessKey'] = attr.text
  85         elif attr.tag == '{https://sts.amazonaws.com/doc/2011-06-15/}SessionToken':
  86             minio_credentials['SessionToken'] = attr.text
  87
  88     return (minio_credentials['AccessKeyId'],
  89             minio_credentials['SecretAccessKey'],
  90             minio_credentials['SessionToken'])
  91
  92
  93 def get_authorization_headers(url, resource):
  94     minio_key, minio_secret, minio_token = get_minio_credentials(url)
  95
  96     date = formatdate(timeval=None, localtime=False, usegmt=True)
  97     to_sign = "GET\n\n\n%s\nx-amz-security-token:%s\n/%s/%s" % (date,
  98                                                                 minio_token,
  99                                                                 OPTIONS.download['minio_bucket'],
 100                                                                 requote_uri(resource))
 101     signature = sign_with_hmac(minio_secret, to_sign)
 102
 103     headers = {'Host': OPTIONS.download['minio_host'],
 104                'Date': date,
 105                'Authorization': 'AWS %s:%s' % (minio_key, signature),
 106                'x-amz-security-token': minio_token}
 107     return headers
 108
 109
 110 def download(url: str, file_path: str, headers: Dict[str, str], attempts: int = 2) -> None:
 111     """Downloads a URL content into a file
 112
 113     :param url: URL to download
 114     :param file_path: Local file name to contain the data downloaded
 115     :param attempts: Number of attempts
 116     """
 117     retries = Retry(
 118         connect=attempts,
 119         read=attempts,
 120         redirect=attempts,
 121         raise_on_redirect=False
 122     )
 123     session = requests.Session()
 124     for protocol in ["http://", "https://"]:
 125         adapter = HTTPAdapter(max_retries=retries)
 126         session.mount(protocol, adapter)
 127     for protocol in ["file://"]:
 128         file_adapter = LocalFileAdapter()
 129         session.mount(protocol, file_adapter)
 130
 131     with session.get(url,
 132                      allow_redirects=True,
 133                      stream=True,
 134                      headers=headers) as req:
 135         with open(file_path, "wb") as file:
 136             req.raise_for_status()
 137             for chunk in req.iter_content(chunk_size=1048576):
 138                 if chunk:
 139                     file.write(chunk)
 140
 141     # Checking file integrity
 142     try:
 143         file_size = int(req.headers["Content-length"])
 144     except KeyError:
 145         print("Error getting Content-Length from server. "
 146               "Skipping file size check.")
 147         return
 148
 149     if file_size != path.getsize(file_path):
 150         raise exceptions.PiglitFatalError(
 151                 f"Invalid filesize src {file_size}"
 152                 f"doesn't match {path.getsize(file_path)}"
 153         )
 154
 155
 156 def ensure_file(file_path):
 157     destination_file_path = path.join(OPTIONS.db_path, file_path)
 158     if OPTIONS.download['url'] is None:
 159         if not path.exists(destination_file_path):
 160             raise exceptions.PiglitFatalError(
 161                 '{} missing'.format(destination_file_path))
 162         return
 163
 164     url = OPTIONS.download['url'].geturl()
 165
 166     if OPTIONS.download['caching_proxy_url'] is not None:
 167         url = OPTIONS.download['caching_proxy_url'].geturl() + url
 168
 169     core.check_dir(path.dirname(destination_file_path))
 170
 171     if not OPTIONS.download['force'] and path.exists(destination_file_path):
 172         return
 173
 174     print('[check_image] Downloading file {}'.format(
 175         file_path), end=' ', flush=True)
 176
 177     if OPTIONS.download['minio_host']:
 178         assert OPTIONS.download['minio_bucket']
 179         assert OPTIONS.download['role_session_name']
 180         assert OPTIONS.download['jwt']
 181         headers = get_authorization_headers(url, file_path)
 182     else:
 183         headers = None
 184
 185     download_time = time()
 186
 187     download(url + file_path, destination_file_path, headers)
 188
 189     print('took %ds.' % (time() - download_time), flush=True)