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