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