1#===----------------------------------------------------------------------===##
2#
3# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4# See https://llvm.org/LICENSE.txt for license information.
5# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6#
7#===----------------------------------------------------------------------===##
8
9from contextlib import contextmanager
10import errno
11import os
12import platform
13import signal
14import subprocess
15import sys
16import tempfile
17import threading
18
19
20# FIXME: Most of these functions are cribbed from LIT
21def to_bytes(str):
22    # Encode to UTF-8 to get binary data.
23    if isinstance(str, bytes):
24        return str
25    return str.encode('utf-8')
26
27def to_string(bytes):
28    if isinstance(bytes, str):
29        return bytes
30    return to_bytes(bytes)
31
32def convert_string(bytes):
33    try:
34        return to_string(bytes.decode('utf-8'))
35    except AttributeError: # 'str' object has no attribute 'decode'.
36        return str(bytes)
37    except UnicodeError:
38        return str(bytes)
39
40
41def cleanFile(filename):
42    try:
43        os.remove(filename)
44    except OSError:
45        pass
46
47
48@contextmanager
49def guardedTempFilename(suffix='', prefix='', dir=None):
50    # Creates and yeilds a temporary filename within a with statement. The file
51    # is removed upon scope exit.
52    handle, name = tempfile.mkstemp(suffix=suffix, prefix=prefix, dir=dir)
53    os.close(handle)
54    yield name
55    cleanFile(name)
56
57
58@contextmanager
59def guardedFilename(name):
60    # yeilds a filename within a with statement. The file is removed upon scope
61    # exit.
62    yield name
63    cleanFile(name)
64
65
66@contextmanager
67def nullContext(value):
68    # yeilds a variable within a with statement. No action is taken upon scope
69    # exit.
70    yield value
71
72
73def makeReport(cmd, out, err, rc):
74    report = "Command: %s\n" % cmd
75    report += "Exit Code: %d\n" % rc
76    if out:
77        report += "Standard Output:\n--\n%s--\n" % out
78    if err:
79        report += "Standard Error:\n--\n%s--\n" % err
80    report += '\n'
81    return report
82
83
84def capture(args, env=None):
85    """capture(command) - Run the given command (or argv list) in a shell and
86    return the standard output. Raises a CalledProcessError if the command
87    exits with a non-zero status."""
88    p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
89                         env=env)
90    out, err = p.communicate()
91    out = convert_string(out)
92    err = convert_string(err)
93    if p.returncode != 0:
94        raise subprocess.CalledProcessError(cmd=args,
95                                            returncode=p.returncode,
96                                            output="{}\n{}".format(out, err))
97    return out
98
99
100def which(command, paths = None):
101    """which(command, [paths]) - Look up the given command in the paths string
102    (or the PATH environment variable, if unspecified)."""
103
104    if paths is None:
105        paths = os.environ.get('PATH','')
106
107    # Check for absolute match first.
108    if os.path.isfile(command):
109        return command
110
111    # Would be nice if Python had a lib function for this.
112    if not paths:
113        paths = os.defpath
114
115    # Get suffixes to search.
116    # On Cygwin, 'PATHEXT' may exist but it should not be used.
117    if os.pathsep == ';':
118        pathext = os.environ.get('PATHEXT', '').split(';')
119    else:
120        pathext = ['']
121
122    # Search the paths...
123    for path in paths.split(os.pathsep):
124        for ext in pathext:
125            p = os.path.join(path, command + ext)
126            if os.path.exists(p) and not os.path.isdir(p):
127                return p
128
129    return None
130
131
132def checkToolsPath(dir, tools):
133    for tool in tools:
134        if not os.path.exists(os.path.join(dir, tool)):
135            return False
136    return True
137
138
139def whichTools(tools, paths):
140    for path in paths.split(os.pathsep):
141        if checkToolsPath(path, tools):
142            return path
143    return None
144
145def mkdir_p(path):
146    """mkdir_p(path) - Make the "path" directory, if it does not exist; this
147    will also make directories for any missing parent directories."""
148    if not path or os.path.exists(path):
149        return
150
151    parent = os.path.dirname(path)
152    if parent != path:
153        mkdir_p(parent)
154
155    try:
156        os.mkdir(path)
157    except OSError:
158        e = sys.exc_info()[1]
159        # Ignore EEXIST, which may occur during a race condition.
160        if e.errno != errno.EEXIST:
161            raise
162
163
164class ExecuteCommandTimeoutException(Exception):
165    def __init__(self, msg, out, err, exitCode, command=None):
166        assert isinstance(msg, str)
167        assert isinstance(out, str)
168        assert isinstance(err, str)
169        assert isinstance(exitCode, int)
170        if command is not None:
171            assert isinstance(command, list) or isinstance(command, str)
172        self.msg = msg
173        self.out = out
174        self.err = err
175        self.exitCode = exitCode
176        self.command = command
177
178# Close extra file handles on UNIX (on Windows this cannot be done while
179# also redirecting input).
180kUseCloseFDs = not (platform.system() == 'Windows')
181def executeCommand(command, cwd=None, env=None, input=None, timeout=0):
182    """
183        Execute command ``command`` (list of arguments or string)
184        with
185        * working directory ``cwd`` (str), use None to use the current
186          working directory
187        * environment ``env`` (dict), use None for none
188        * Input to the command ``input`` (str), use string to pass
189          no input.
190        * Max execution time ``timeout`` (int) seconds. Use 0 for no timeout.
191
192        Returns a tuple (out, err, exitCode) where
193        * ``out`` (str) is the standard output of running the command
194        * ``err`` (str) is the standard error of running the command
195        * ``exitCode`` (int) is the exitCode of running the command
196
197        If the timeout is hit an ``ExecuteCommandTimeoutException``
198        is raised.
199    """
200    if input is not None:
201        input = to_bytes(input)
202    p = subprocess.Popen(command, cwd=cwd,
203                         stdin=subprocess.PIPE,
204                         stdout=subprocess.PIPE,
205                         stderr=subprocess.PIPE,
206                         env=env, close_fds=kUseCloseFDs)
207    timerObject = None
208    # FIXME: Because of the way nested function scopes work in Python 2.x we
209    # need to use a reference to a mutable object rather than a plain
210    # bool. In Python 3 we could use the "nonlocal" keyword but we need
211    # to support Python 2 as well.
212    hitTimeOut = [False]
213    # For python >= 3.3 we can use the timeout parameter to p.communicate which
214    # is much less fragile than the timer code.
215    if sys.version_info > (3, 3):
216        try:
217            if timeout <= 0:
218                timeout = None
219            out, err = p.communicate(input=input, timeout=timeout)
220        except subprocess.TimeoutExpired:
221            killProcessAndChildren(p.pid)
222            # p.communicate should return immediately and return the remaining
223            # stdout/stderr until the process was killed.
224            out, err = p.communicate(input=input)
225            hitTimeOut[0] = True
226        finally:
227            exitCode = p.returncode
228    else:  # python < 3.3
229        try:
230            if timeout > 0:
231                def killProcess():
232                    # We may be invoking a shell so we need to kill the
233                    # process and all its children.
234                    hitTimeOut[0] = True
235                    killProcessAndChildren(p.pid)
236
237                timerObject = threading.Timer(timeout, killProcess)
238                timerObject.start()
239
240            kwargs = dict()
241            out, err = p.communicate(input=input)
242            exitCode = p.wait()
243        finally:
244            if timerObject != None:
245                timerObject.cancel()
246
247    # Ensure the resulting output is always of string type.
248    out = convert_string(out)
249    err = convert_string(err)
250
251    if hitTimeOut[0]:
252        raise ExecuteCommandTimeoutException(
253            msg='Reached timeout of {} seconds'.format(timeout),
254            out=out,
255            err=err,
256            exitCode=exitCode,
257            command=command
258            )
259
260    # Detect Ctrl-C in subprocess.
261    if exitCode == -signal.SIGINT:
262        raise KeyboardInterrupt
263
264    return out, err, exitCode
265
266
267def killProcessAndChildren(pid):
268    """
269    This function kills a process with ``pid`` and all its
270    running children (recursively). It is currently implemented
271    using the psutil module which provides a simple platform
272    neutral implementation.
273
274    TODO: Reimplement this without using psutil so we can
275          remove our dependency on it.
276    """
277    if platform.system() == 'AIX':
278        subprocess.call('kill -kill $(ps -o pid= -L{})'.format(pid), shell=True)
279    else:
280        import psutil
281        try:
282            psutilProc = psutil.Process(pid)
283            # Handle the different psutil API versions
284            try:
285                # psutil >= 2.x
286                children_iterator = psutilProc.children(recursive=True)
287            except AttributeError:
288                # psutil 1.x
289                children_iterator = psutilProc.get_children(recursive=True)
290            for child in children_iterator:
291                try:
292                    child.kill()
293                except psutil.NoSuchProcess:
294                    pass
295            psutilProc.kill()
296        except psutil.NoSuchProcess:
297            pass
298
299
300def executeCommandVerbose(cmd, *args, **kwargs):
301    """
302    Execute a command and print its output on failure.
303    """
304    out, err, exitCode = executeCommand(cmd, *args, **kwargs)
305    if exitCode != 0:
306        report = makeReport(cmd, out, err, exitCode)
307        report += "\n\nFailed!"
308        sys.stderr.write('%s\n' % report)
309    return out, err, exitCode
310
311
312def executeCommandOrDie(cmd, *args, **kwargs):
313    """
314    Execute a command and print its output on failure.
315    """
316    out, err, exitCode = executeCommand(cmd, *args, **kwargs)
317    if exitCode != 0:
318        report = makeReport(cmd, out, err, exitCode)
319        report += "\n\nFailed!"
320        sys.stderr.write('%s\n' % report)
321        sys.exit(exitCode)
322    return out, err, exitCode
323