1"""Utility functions for Raptor""" 2# This Source Code Form is subject to the terms of the Mozilla Public 3# License, v. 2.0. If a copy of the MPL was not distributed with this 4# file, You can obtain one at http://mozilla.org/MPL/2.0/. 5 6from __future__ import absolute_import 7 8import subprocess 9import time 10import bz2 11import gzip 12import os 13import signal 14import sys 15import socket 16 17from six.moves.urllib.request import urlretrieve 18from redo import retriable, retry 19 20try: 21 import zstandard 22except ImportError: 23 zstandard = None 24try: 25 import lzma 26except ImportError: 27 lzma = None 28 29from mozlog import get_proxy_logger 30from mozprocess import ProcessHandler 31from mozproxy import mozharness_dir, mozbase_dir 32 33 34LOG = get_proxy_logger(component="mozproxy") 35 36# running locally via mach 37TOOLTOOL_PATHS = [ 38 os.path.join(mozharness_dir, "external_tools", "tooltool.py"), 39 os.path.join( 40 mozbase_dir, 41 "..", 42 "..", 43 "python", 44 "mozbuild", 45 "mozbuild", 46 "action", 47 "tooltool.py", 48 ), 49] 50 51if "MOZ_UPLOAD_DIR" in os.environ: 52 TOOLTOOL_PATHS.append( 53 os.path.join( 54 os.environ["MOZ_UPLOAD_DIR"], 55 "..", 56 "..", 57 "mozharness", 58 "external_tools", 59 "tooltool.py", 60 ) 61 ) 62 63 64def transform_platform(str_to_transform, config_platform, config_processor=None): 65 """Transform platform name i.e. 'mitmproxy-rel-bin-{platform}.manifest' 66 67 transforms to 'mitmproxy-rel-bin-osx.manifest'. 68 Also transform '{x64}' if needed for 64 bit / win 10 69 """ 70 if "{platform}" not in str_to_transform and "{x64}" not in str_to_transform: 71 return str_to_transform 72 73 if "win" in config_platform: 74 platform_id = "win" 75 elif config_platform == "mac": 76 platform_id = "osx" 77 else: 78 platform_id = "linux64" 79 80 if "{platform}" in str_to_transform: 81 str_to_transform = str_to_transform.replace("{platform}", platform_id) 82 83 if "{x64}" in str_to_transform and config_processor is not None: 84 if "x86_64" in config_processor: 85 str_to_transform = str_to_transform.replace("{x64}", "_x64") 86 else: 87 str_to_transform = str_to_transform.replace("{x64}", "") 88 89 return str_to_transform 90 91 92@retriable(sleeptime=2) 93def tooltool_download(manifest, run_local, raptor_dir): 94 """Download a file from tooltool using the provided tooltool manifest""" 95 96 def outputHandler(line): 97 LOG.info(line) 98 99 tooltool_path = None 100 101 for path in TOOLTOOL_PATHS: 102 if os.path.exists(os.path.dirname(path)): 103 tooltool_path = path 104 break 105 if tooltool_path is None: 106 raise Exception("Could not find tooltool path!") 107 108 if run_local: 109 command = [sys.executable, tooltool_path, "fetch", "-o", "-m", manifest] 110 else: 111 # Attempt to determine the tooltool cache path: 112 # - TOOLTOOLCACHE is used by Raptor tests 113 # - TOOLTOOL_CACHE is automatically set for taskcluster jobs 114 # - fallback to a hardcoded path 115 _cache = next( 116 x 117 for x in ( 118 os.environ.get("TOOLTOOLCACHE"), 119 os.environ.get("TOOLTOOL_CACHE"), 120 "/builds/tooltool_cache", 121 ) 122 if x is not None 123 ) 124 125 command = [ 126 sys.executable, 127 tooltool_path, 128 "fetch", 129 "-o", 130 "-m", 131 manifest, 132 "-c", 133 _cache, 134 ] 135 136 try: 137 proc = ProcessHandler( 138 command, processOutputLine=outputHandler, storeOutput=False, cwd=raptor_dir 139 ) 140 proc.run() 141 if proc.wait() != 0: 142 raise Exception("Command failed") 143 except Exception as e: 144 LOG.critical( 145 "Error while downloading {} from tooltool:{}".format(manifest, str(e)) 146 ) 147 if proc.poll() is None: 148 proc.kill(signal.SIGTERM) 149 raise 150 151 152def archive_type(path): 153 filename, extension = os.path.splitext(path) 154 filename, extension2 = os.path.splitext(filename) 155 if extension2 != "": 156 extension = extension2 157 if extension == ".tar": 158 return "tar" 159 elif extension == ".zip": 160 return "zip" 161 return None 162 163 164def extract_archive(path, dest_dir, typ): 165 """Extract an archive to a destination directory.""" 166 167 # Resolve paths to absolute variants. 168 path = os.path.abspath(path) 169 dest_dir = os.path.abspath(dest_dir) 170 suffix = os.path.splitext(path)[-1] 171 172 # We pipe input to the decompressor program so that we can apply 173 # custom decompressors that the program may not know about. 174 if typ == "tar": 175 if suffix == ".bz2": 176 ifh = bz2.open(str(path), "rb") 177 elif suffix == ".gz": 178 ifh = gzip.open(str(path), "rb") 179 elif suffix == ".xz": 180 if not lzma: 181 raise ValueError("lzma Python package not available") 182 ifh = lzma.open(str(path), "rb") 183 elif suffix == ".zst": 184 if not zstandard: 185 raise ValueError("zstandard Python package not available") 186 dctx = zstandard.ZstdDecompressor() 187 ifh = dctx.stream_reader(path.open("rb")) 188 elif suffix == ".tar": 189 ifh = path.open("rb") 190 else: 191 raise ValueError("unknown archive format for tar file: %s" % path) 192 args = ["tar", "xf", "-"] 193 pipe_stdin = True 194 elif typ == "zip": 195 # unzip from stdin has wonky behavior. We don't use a pipe for it. 196 ifh = open(os.devnull, "rb") 197 args = ["unzip", "-o", str(path)] 198 pipe_stdin = False 199 else: 200 raise ValueError("unknown archive format: %s" % path) 201 202 LOG.info("Extracting %s to %s using %r" % (path, dest_dir, args)) 203 t0 = time.time() 204 with ifh: 205 p = subprocess.Popen(args, cwd=str(dest_dir), bufsize=0, stdin=subprocess.PIPE) 206 while True: 207 if not pipe_stdin: 208 break 209 chunk = ifh.read(131072) 210 if not chunk: 211 break 212 p.stdin.write(chunk) 213 # make sure we wait for the command to finish 214 p.communicate() 215 216 if p.returncode: 217 raise Exception("%r exited %d" % (args, p.returncode)) 218 LOG.info("%s extracted in %.3fs" % (path, time.time() - t0)) 219 220 221def download_file_from_url(url, local_dest, extract=False): 222 """Receive a file in a URL and download it, i.e. for the hostutils tooltool manifest 223 the url received would be formatted like this: 224 config/tooltool-manifests/linux64/hostutils.manifest""" 225 if os.path.exists(local_dest): 226 LOG.info("file already exists at: %s" % local_dest) 227 if not extract: 228 return True 229 else: 230 LOG.info("downloading: %s to %s" % (url, local_dest)) 231 try: 232 retry(urlretrieve, args=(url, local_dest), attempts=3, sleeptime=5) 233 except Exception: 234 LOG.error("Failed to download file: %s" % local_dest, exc_info=True) 235 if os.path.exists(local_dest): 236 # delete partial downloaded file 237 os.remove(local_dest) 238 return False 239 240 if not extract: 241 return os.path.exists(local_dest) 242 243 typ = archive_type(local_dest) 244 if typ is None: 245 LOG.info("Not able to determine archive type for: %s" % local_dest) 246 return False 247 248 extract_archive(local_dest, os.path.dirname(local_dest), typ) 249 return True 250 251 252def get_available_port(): 253 s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 254 s.bind(("", 0)) 255 s.listen(1) 256 port = s.getsockname()[1] 257 s.close() 258 return port 259