1# -*- coding:utf-8 -*-
2# Copyright 2016 The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16"""Various utility functions."""
17
18from __future__ import print_function
19
20import errno
21import functools
22import os
23import signal
24import subprocess
25import sys
26import tempfile
27import time
28
29_path = os.path.realpath(__file__ + '/../..')
30if sys.path[0] != _path:
31    sys.path.insert(0, _path)
32del _path
33
34# pylint: disable=wrong-import-position
35import rh.shell
36import rh.signals
37from rh.sixish import string_types
38
39
40def timedelta_str(delta):
41    """A less noisy timedelta.__str__.
42
43    The default timedelta stringification contains a lot of leading zeros and
44    uses microsecond resolution.  This makes for noisy output.
45    """
46    total = delta.total_seconds()
47    hours, rem = divmod(total, 3600)
48    mins, secs = divmod(rem, 60)
49    ret = '%i.%03is' % (secs, delta.microseconds // 1000)
50    if mins:
51        ret = '%im%s' % (mins, ret)
52    if hours:
53        ret = '%ih%s' % (hours, ret)
54    return ret
55
56
57class CompletedProcess(getattr(subprocess, 'CompletedProcess', object)):
58    """An object to store various attributes of a child process.
59
60    This is akin to subprocess.CompletedProcess.
61    """
62
63    # The linter is confused by the getattr usage above.
64    # TODO(vapier): Drop this once we're Python 3-only and we drop getattr.
65    # pylint: disable=bad-option-value,super-on-old-class
66    def __init__(self, args=None, returncode=None, stdout=None, stderr=None):
67        if sys.version_info.major < 3:
68            self.args = args
69            self.stdout = stdout
70            self.stderr = stderr
71            self.returncode = returncode
72        else:
73            super(CompletedProcess, self).__init__(
74                args=args, returncode=returncode, stdout=stdout, stderr=stderr)
75
76    @property
77    def cmd(self):
78        """Alias to self.args to better match other subprocess APIs."""
79        return self.args
80
81    @property
82    def cmdstr(self):
83        """Return self.cmd as a nicely formatted string (useful for logs)."""
84        return rh.shell.cmd_to_str(self.cmd)
85
86
87class CalledProcessError(subprocess.CalledProcessError):
88    """Error caught in run() function.
89
90    This is akin to subprocess.CalledProcessError.  We do not support |output|,
91    only |stdout|.
92
93    Attributes:
94      returncode: The exit code of the process.
95      cmd: The command that triggered this exception.
96      msg: Short explanation of the error.
97      exception: The underlying Exception if available.
98    """
99
100    def __init__(self, returncode, cmd, stdout=None, stderr=None, msg=None,
101                 exception=None):
102        if exception is not None and not isinstance(exception, Exception):
103            raise TypeError('exception must be an exception instance; got %r'
104                            % (exception,))
105
106        super(CalledProcessError, self).__init__(returncode, cmd, stdout)
107        # The parent class will set |output|, so delete it.
108        del self.output
109        # TODO(vapier): When we're Python 3-only, delete this assignment as the
110        # parent handles it for us.
111        self.stdout = stdout
112        # TODO(vapier): When we're Python 3-only, move stderr to the init above.
113        self.stderr = stderr
114        self.msg = msg
115        self.exception = exception
116
117    @property
118    def cmdstr(self):
119        """Return self.cmd as a well shell-quoted string for debugging."""
120        return '' if self.cmd is None else rh.shell.cmd_to_str(self.cmd)
121
122    def stringify(self, stdout=True, stderr=True):
123        """Custom method for controlling what is included in stringifying this.
124
125        Args:
126          stdout: Whether to include captured stdout in the return value.
127          stderr: Whether to include captured stderr in the return value.
128
129        Returns:
130          A summary string for this result.
131        """
132        items = [
133            'return code: %s; command: %s' % (self.returncode, self.cmdstr),
134        ]
135        if stderr and self.stderr:
136            items.append(self.stderr)
137        if stdout and self.stdout:
138            items.append(self.stdout)
139        if self.msg:
140            items.append(self.msg)
141        return '\n'.join(items)
142
143    def __str__(self):
144        return self.stringify()
145
146
147class TerminateCalledProcessError(CalledProcessError):
148    """We were signaled to shutdown while running a command.
149
150    Client code shouldn't generally know, nor care about this class.  It's
151    used internally to suppress retry attempts when we're signaled to die.
152    """
153
154
155def _kill_child_process(proc, int_timeout, kill_timeout, cmd, original_handler,
156                        signum, frame):
157    """Used as a signal handler by RunCommand.
158
159    This is internal to Runcommand.  No other code should use this.
160    """
161    if signum:
162        # If we've been invoked because of a signal, ignore delivery of that
163        # signal from this point forward.  The invoking context of this func
164        # restores signal delivery to what it was prior; we suppress future
165        # delivery till then since this code handles SIGINT/SIGTERM fully
166        # including delivering the signal to the original handler on the way
167        # out.
168        signal.signal(signum, signal.SIG_IGN)
169
170    # Do not trust Popen's returncode alone; we can be invoked from contexts
171    # where the Popen instance was created, but no process was generated.
172    if proc.returncode is None and proc.pid is not None:
173        try:
174            while proc.poll_lock_breaker() is None and int_timeout >= 0:
175                time.sleep(0.1)
176                int_timeout -= 0.1
177
178            proc.terminate()
179            while proc.poll_lock_breaker() is None and kill_timeout >= 0:
180                time.sleep(0.1)
181                kill_timeout -= 0.1
182
183            if proc.poll_lock_breaker() is None:
184                # Still doesn't want to die.  Too bad, so sad, time to die.
185                proc.kill()
186        except EnvironmentError as e:
187            print('Ignoring unhandled exception in _kill_child_process: %s' % e,
188                  file=sys.stderr)
189
190        # Ensure our child process has been reaped.
191        kwargs = {}
192        if sys.version_info.major >= 3:
193            # ... but don't wait forever.
194            kwargs['timeout'] = 60
195        proc.wait_lock_breaker(**kwargs)
196
197    if not rh.signals.relay_signal(original_handler, signum, frame):
198        # Mock up our own, matching exit code for signaling.
199        raise TerminateCalledProcessError(
200            signum << 8, cmd, msg='Received signal %i' % signum)
201
202
203class _Popen(subprocess.Popen):
204    """subprocess.Popen derivative customized for our usage.
205
206    Specifically, we fix terminate/send_signal/kill to work if the child process
207    was a setuid binary; on vanilla kernels, the parent can wax the child
208    regardless, on goobuntu this apparently isn't allowed, thus we fall back
209    to the sudo machinery we have.
210
211    While we're overriding send_signal, we also suppress ESRCH being raised
212    if the process has exited, and suppress signaling all together if the
213    process has knowingly been waitpid'd already.
214    """
215
216    # pylint: disable=arguments-differ
217    def send_signal(self, signum):
218        if self.returncode is not None:
219            # The original implementation in Popen allows signaling whatever
220            # process now occupies this pid, even if the Popen object had
221            # waitpid'd.  Since we can escalate to sudo kill, we do not want
222            # to allow that.  Fixing this addresses that angle, and makes the
223            # API less sucky in the process.
224            return
225
226        try:
227            os.kill(self.pid, signum)
228        except EnvironmentError as e:
229            if e.errno == errno.ESRCH:
230                # Since we know the process is dead, reap it now.
231                # Normally Popen would throw this error- we suppress it since
232                # frankly that's a misfeature and we're already overriding
233                # this method.
234                self.poll()
235            else:
236                raise
237
238    def _lock_breaker(self, func, *args, **kwargs):
239        """Helper to manage the waitpid lock.
240
241        Workaround https://bugs.python.org/issue25960.
242        """
243        # If the lock doesn't exist, or is not locked, call the func directly.
244        lock = getattr(self, '_waitpid_lock', None)
245        if lock is not None and lock.locked():
246            try:
247                lock.release()
248                return func(*args, **kwargs)
249            finally:
250                if not lock.locked():
251                    lock.acquire()
252        else:
253            return func(*args, **kwargs)
254
255    def poll_lock_breaker(self, *args, **kwargs):
256        """Wrapper around poll() to break locks if needed."""
257        return self._lock_breaker(self.poll, *args, **kwargs)
258
259    def wait_lock_breaker(self, *args, **kwargs):
260        """Wrapper around wait() to break locks if needed."""
261        return self._lock_breaker(self.wait, *args, **kwargs)
262
263
264# We use the keyword arg |input| which trips up pylint checks.
265# pylint: disable=redefined-builtin,input-builtin
266def run(cmd, redirect_stdout=False, redirect_stderr=False, cwd=None, input=None,
267        shell=False, env=None, extra_env=None, combine_stdout_stderr=False,
268        check=True, int_timeout=1, kill_timeout=1, capture_output=False,
269        close_fds=True):
270    """Runs a command.
271
272    Args:
273      cmd: cmd to run.  Should be input to subprocess.Popen.  If a string, shell
274          must be true.  Otherwise the command must be an array of arguments,
275          and shell must be false.
276      redirect_stdout: Returns the stdout.
277      redirect_stderr: Holds stderr output until input is communicated.
278      cwd: The working directory to run this cmd.
279      input: The data to pipe into this command through stdin.  If a file object
280          or file descriptor, stdin will be connected directly to that.
281      shell: Controls whether we add a shell as a command interpreter.  See cmd
282          since it has to agree as to the type.
283      env: If non-None, this is the environment for the new process.
284      extra_env: If set, this is added to the environment for the new process.
285          This dictionary is not used to clear any entries though.
286      combine_stdout_stderr: Combines stdout and stderr streams into stdout.
287      check: Whether to raise an exception when command returns a non-zero exit
288          code, or return the CompletedProcess object containing the exit code.
289          Note: will still raise an exception if the cmd file does not exist.
290      int_timeout: If we're interrupted, how long (in seconds) should we give
291          the invoked process to clean up before we send a SIGTERM.
292      kill_timeout: If we're interrupted, how long (in seconds) should we give
293          the invoked process to shutdown from a SIGTERM before we SIGKILL it.
294      capture_output: Set |redirect_stdout| and |redirect_stderr| to True.
295      close_fds: Whether to close all fds before running |cmd|.
296
297    Returns:
298      A CompletedProcess object.
299
300    Raises:
301      CalledProcessError: Raises exception on error.
302    """
303    if capture_output:
304        redirect_stdout, redirect_stderr = True, True
305
306    # Set default for variables.
307    popen_stdout = None
308    popen_stderr = None
309    stdin = None
310    result = CompletedProcess()
311
312    # Force the timeout to float; in the process, if it's not convertible,
313    # a self-explanatory exception will be thrown.
314    kill_timeout = float(kill_timeout)
315
316    def _get_tempfile():
317        kwargs = {}
318        if sys.version_info.major < 3:
319            kwargs['bufsize'] = 0
320        else:
321            kwargs['buffering'] = 0
322        try:
323            return tempfile.TemporaryFile(**kwargs)
324        except EnvironmentError as e:
325            if e.errno != errno.ENOENT:
326                raise
327            # This can occur if we were pointed at a specific location for our
328            # TMP, but that location has since been deleted.  Suppress that
329            # issue in this particular case since our usage gurantees deletion,
330            # and since this is primarily triggered during hard cgroups
331            # shutdown.
332            return tempfile.TemporaryFile(dir='/tmp', **kwargs)
333
334    # Modify defaults based on parameters.
335    # Note that tempfiles must be unbuffered else attempts to read
336    # what a separate process did to that file can result in a bad
337    # view of the file.
338    # The Popen API accepts either an int or a file handle for stdout/stderr.
339    # pylint: disable=redefined-variable-type
340    if redirect_stdout:
341        popen_stdout = _get_tempfile()
342
343    if combine_stdout_stderr:
344        popen_stderr = subprocess.STDOUT
345    elif redirect_stderr:
346        popen_stderr = _get_tempfile()
347    # pylint: enable=redefined-variable-type
348
349    # If subprocesses have direct access to stdout or stderr, they can bypass
350    # our buffers, so we need to flush to ensure that output is not interleaved.
351    if popen_stdout is None or popen_stderr is None:
352        sys.stdout.flush()
353        sys.stderr.flush()
354
355    # If input is a string, we'll create a pipe and send it through that.
356    # Otherwise we assume it's a file object that can be read from directly.
357    if isinstance(input, string_types):
358        stdin = subprocess.PIPE
359        input = input.encode('utf-8')
360    elif input is not None:
361        stdin = input
362        input = None
363
364    if isinstance(cmd, string_types):
365        if not shell:
366            raise Exception('Cannot run a string command without a shell')
367        cmd = ['/bin/bash', '-c', cmd]
368        shell = False
369    elif shell:
370        raise Exception('Cannot run an array command with a shell')
371
372    # If we are using enter_chroot we need to use enterchroot pass env through
373    # to the final command.
374    env = env.copy() if env is not None else os.environ.copy()
375    env.update(extra_env if extra_env else {})
376
377    result.args = cmd
378
379    proc = None
380    try:
381        proc = _Popen(cmd, cwd=cwd, stdin=stdin, stdout=popen_stdout,
382                      stderr=popen_stderr, shell=False, env=env,
383                      close_fds=close_fds)
384
385        old_sigint = signal.getsignal(signal.SIGINT)
386        handler = functools.partial(_kill_child_process, proc, int_timeout,
387                                    kill_timeout, cmd, old_sigint)
388        signal.signal(signal.SIGINT, handler)
389
390        old_sigterm = signal.getsignal(signal.SIGTERM)
391        handler = functools.partial(_kill_child_process, proc, int_timeout,
392                                    kill_timeout, cmd, old_sigterm)
393        signal.signal(signal.SIGTERM, handler)
394
395        try:
396            (result.stdout, result.stderr) = proc.communicate(input)
397        finally:
398            signal.signal(signal.SIGINT, old_sigint)
399            signal.signal(signal.SIGTERM, old_sigterm)
400
401            if popen_stdout:
402                # The linter is confused by how stdout is a file & an int.
403                # pylint: disable=maybe-no-member,no-member
404                popen_stdout.seek(0)
405                result.stdout = popen_stdout.read()
406                popen_stdout.close()
407
408            if popen_stderr and popen_stderr != subprocess.STDOUT:
409                # The linter is confused by how stderr is a file & an int.
410                # pylint: disable=maybe-no-member,no-member
411                popen_stderr.seek(0)
412                result.stderr = popen_stderr.read()
413                popen_stderr.close()
414
415        result.returncode = proc.returncode
416
417        if check and proc.returncode:
418            msg = 'cwd=%s' % cwd
419            if extra_env:
420                msg += ', extra env=%s' % extra_env
421            raise CalledProcessError(
422                result.returncode, result.cmd, stdout=result.stdout,
423                stderr=result.stderr, msg=msg)
424    except OSError as e:
425        estr = str(e)
426        if e.errno == errno.EACCES:
427            estr += '; does the program need `chmod a+x`?'
428        if not check:
429            result = CompletedProcess(
430                args=cmd, stderr=estr.encode('utf-8'), returncode=255)
431        else:
432            raise CalledProcessError(
433                result.returncode, result.cmd, stdout=result.stdout,
434                stderr=result.stderr, msg=estr, exception=e)
435    finally:
436        if proc is not None:
437            # Ensure the process is dead.
438            # Some pylint3 versions are confused here.
439            # pylint: disable=too-many-function-args
440            _kill_child_process(proc, int_timeout, kill_timeout, cmd, None,
441                                None, None)
442
443    # Make sure output is returned as a string rather than bytes.
444    if result.stdout is not None:
445        result.stdout = result.stdout.decode('utf-8', 'replace')
446    if result.stderr is not None:
447        result.stderr = result.stderr.decode('utf-8', 'replace')
448
449    return result
450# pylint: enable=redefined-builtin,input-builtin
451