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
37
38
39class CommandResult(object):
40    """An object to store various attributes of a child process."""
41
42    def __init__(self, cmd=None, error=None, output=None, returncode=None):
43        self.cmd = cmd
44        self.error = error
45        self.output = output
46        self.returncode = returncode
47
48    @property
49    def cmdstr(self):
50        """Return self.cmd as a nicely formatted string (useful for logs)."""
51        return rh.shell.cmd_to_str(self.cmd)
52
53
54class RunCommandError(Exception):
55    """Error caught in RunCommand() method."""
56
57    def __init__(self, msg, result, exception=None):
58        self.msg, self.result, self.exception = msg, result, exception
59        if exception is not None and not isinstance(exception, Exception):
60            raise ValueError('exception must be an exception instance; got %r'
61                             % (exception,))
62        Exception.__init__(self, msg)
63        self.args = (msg, result, exception)
64
65    def stringify(self, error=True, output=True):
66        """Custom method for controlling what is included in stringifying this.
67
68        Each individual argument is the literal name of an attribute
69        on the result object; if False, that value is ignored for adding
70        to this string content.  If true, it'll be incorporated.
71
72        Args:
73          error: See comment about individual arguments above.
74          output: See comment about individual arguments above.
75        """
76        items = [
77            'return code: %s; command: %s' % (
78                self.result.returncode, self.result.cmdstr),
79        ]
80        if error and self.result.error:
81            items.append(self.result.error)
82        if output and self.result.output:
83            items.append(self.result.output)
84        if self.msg:
85            items.append(self.msg)
86        return '\n'.join(items)
87
88    def __str__(self):
89        # __str__ needs to return ascii, thus force a conversion to be safe.
90        return self.stringify().decode('utf-8', 'replace').encode(
91            'ascii', 'xmlcharrefreplace')
92
93    def __eq__(self, other):
94        return (type(self) == type(other) and
95                self.args == other.args)
96
97    def __ne__(self, other):
98        return not self.__eq__(other)
99
100
101class TerminateRunCommandError(RunCommandError):
102    """We were signaled to shutdown while running a command.
103
104    Client code shouldn't generally know, nor care about this class.  It's
105    used internally to suppress retry attempts when we're signaled to die.
106    """
107
108
109def sudo_run_command(cmd, user='root', **kwargs):
110    """Run a command via sudo.
111
112    Client code must use this rather than coming up with their own RunCommand
113    invocation that jams sudo in- this function is used to enforce certain
114    rules in our code about sudo usage, and as a potential auditing point.
115
116    Args:
117      cmd: The command to run.  See RunCommand for rules of this argument-
118          SudoRunCommand purely prefixes it with sudo.
119      user: The user to run the command as.
120      kwargs: See RunCommand options, it's a direct pass thru to it.
121          Note that this supports a 'strict' keyword that defaults to True.
122          If set to False, it'll suppress strict sudo behavior.
123
124    Returns:
125      See RunCommand documentation.
126
127    Raises:
128      This function may immediately raise RunCommandError if we're operating
129      in a strict sudo context and the API is being misused.
130      Barring that, see RunCommand's documentation- it can raise the same things
131      RunCommand does.
132    """
133    sudo_cmd = ['sudo']
134
135    if user == 'root' and os.geteuid() == 0:
136        return run_command(cmd, **kwargs)
137
138    if user != 'root':
139        sudo_cmd += ['-u', user]
140
141    # Pass these values down into the sudo environment, since sudo will
142    # just strip them normally.
143    extra_env = kwargs.pop('extra_env', None)
144    extra_env = {} if extra_env is None else extra_env.copy()
145
146    sudo_cmd.extend('%s=%s' % (k, v) for k, v in extra_env.iteritems())
147
148    # Finally, block people from passing options to sudo.
149    sudo_cmd.append('--')
150
151    if isinstance(cmd, basestring):
152        # We need to handle shell ourselves so the order is correct:
153        #  $ sudo [sudo args] -- bash -c '[shell command]'
154        # If we let RunCommand take care of it, we'd end up with:
155        #  $ bash -c 'sudo [sudo args] -- [shell command]'
156        shell = kwargs.pop('shell', False)
157        if not shell:
158            raise Exception('Cannot run a string command without a shell')
159        sudo_cmd.extend(['/bin/bash', '-c', cmd])
160    else:
161        sudo_cmd.extend(cmd)
162
163    return run_command(sudo_cmd, **kwargs)
164
165
166def _kill_child_process(proc, int_timeout, kill_timeout, cmd, original_handler,
167                        signum, frame):
168    """Used as a signal handler by RunCommand.
169
170    This is internal to Runcommand.  No other code should use this.
171    """
172    if signum:
173        # If we've been invoked because of a signal, ignore delivery of that
174        # signal from this point forward.  The invoking context of this func
175        # restores signal delivery to what it was prior; we suppress future
176        # delivery till then since this code handles SIGINT/SIGTERM fully
177        # including delivering the signal to the original handler on the way
178        # out.
179        signal.signal(signum, signal.SIG_IGN)
180
181    # Do not trust Popen's returncode alone; we can be invoked from contexts
182    # where the Popen instance was created, but no process was generated.
183    if proc.returncode is None and proc.pid is not None:
184        try:
185            while proc.poll() is None and int_timeout >= 0:
186                time.sleep(0.1)
187                int_timeout -= 0.1
188
189            proc.terminate()
190            while proc.poll() is None and kill_timeout >= 0:
191                time.sleep(0.1)
192                kill_timeout -= 0.1
193
194            if proc.poll() is None:
195                # Still doesn't want to die.  Too bad, so sad, time to die.
196                proc.kill()
197        except EnvironmentError as e:
198            print('Ignoring unhandled exception in _kill_child_process: %s' % e,
199                  file=sys.stderr)
200
201        # Ensure our child process has been reaped.
202        proc.wait()
203
204    if not rh.signals.relay_signal(original_handler, signum, frame):
205        # Mock up our own, matching exit code for signaling.
206        cmd_result = CommandResult(cmd=cmd, returncode=signum << 8)
207        raise TerminateRunCommandError('Received signal %i' % signum,
208                                       cmd_result)
209
210
211class _Popen(subprocess.Popen):
212    """subprocess.Popen derivative customized for our usage.
213
214    Specifically, we fix terminate/send_signal/kill to work if the child process
215    was a setuid binary; on vanilla kernels, the parent can wax the child
216    regardless, on goobuntu this apparently isn't allowed, thus we fall back
217    to the sudo machinery we have.
218
219    While we're overriding send_signal, we also suppress ESRCH being raised
220    if the process has exited, and suppress signaling all together if the
221    process has knowingly been waitpid'd already.
222    """
223
224    def send_signal(self, signum):
225        if self.returncode is not None:
226            # The original implementation in Popen allows signaling whatever
227            # process now occupies this pid, even if the Popen object had
228            # waitpid'd.  Since we can escalate to sudo kill, we do not want
229            # to allow that.  Fixing this addresses that angle, and makes the
230            # API less sucky in the process.
231            return
232
233        try:
234            os.kill(self.pid, signum)
235        except EnvironmentError as e:
236            if e.errno == errno.EPERM:
237                # Kill returns either 0 (signal delivered), or 1 (signal wasn't
238                # delivered).  This isn't particularly informative, but we still
239                # need that info to decide what to do, thus error_code_ok=True.
240                ret = sudo_run_command(['kill', '-%i' % signum, str(self.pid)],
241                                       redirect_stdout=True,
242                                       redirect_stderr=True, error_code_ok=True)
243                if ret.returncode == 1:
244                    # The kill binary doesn't distinguish between permission
245                    # denied and the pid is missing.  Denied can only occur
246                    # under weird grsec/selinux policies.  We ignore that
247                    # potential and just assume the pid was already dead and
248                    # try to reap it.
249                    self.poll()
250            elif e.errno == errno.ESRCH:
251                # Since we know the process is dead, reap it now.
252                # Normally Popen would throw this error- we suppress it since
253                # frankly that's a misfeature and we're already overriding
254                # this method.
255                self.poll()
256            else:
257                raise
258
259
260# pylint: disable=redefined-builtin
261def run_command(cmd, error_message=None, redirect_stdout=False,
262                redirect_stderr=False, cwd=None, input=None,
263                shell=False, env=None, extra_env=None, ignore_sigint=False,
264                combine_stdout_stderr=False, log_stdout_to_file=None,
265                error_code_ok=False, int_timeout=1, kill_timeout=1,
266                stdout_to_pipe=False, capture_output=False,
267                quiet=False, close_fds=True):
268    """Runs a command.
269
270    Args:
271      cmd: cmd to run.  Should be input to subprocess.Popen.  If a string, shell
272          must be true.  Otherwise the command must be an array of arguments,
273          and shell must be false.
274      error_message: Prints out this message when an error occurs.
275      redirect_stdout: Returns the stdout.
276      redirect_stderr: Holds stderr output until input is communicated.
277      cwd: The working directory to run this cmd.
278      input: The data to pipe into this command through stdin.  If a file object
279          or file descriptor, stdin will be connected directly to that.
280      shell: Controls whether we add a shell as a command interpreter.  See cmd
281          since it has to agree as to the type.
282      env: If non-None, this is the environment for the new process.
283      extra_env: If set, this is added to the environment for the new process.
284          This dictionary is not used to clear any entries though.
285      ignore_sigint: If True, we'll ignore signal.SIGINT before calling the
286          child.  This is the desired behavior if we know our child will handle
287          Ctrl-C.  If we don't do this, I think we and the child will both get
288          Ctrl-C at the same time, which means we'll forcefully kill the child.
289      combine_stdout_stderr: Combines stdout and stderr streams into stdout.
290      log_stdout_to_file: If set, redirects stdout to file specified by this
291          path.  If |combine_stdout_stderr| is set to True, then stderr will
292          also be logged to the specified file.
293      error_code_ok: Does not raise an exception when command returns a non-zero
294          exit code.  Instead, returns the CommandResult object containing the
295          exit code.
296      int_timeout: If we're interrupted, how long (in seconds) should we give
297          the invoked process to clean up before we send a SIGTERM.
298      kill_timeout: If we're interrupted, how long (in seconds) should we give
299          the invoked process to shutdown from a SIGTERM before we SIGKILL it.
300      stdout_to_pipe: Redirect stdout to pipe.
301      capture_output: Set |redirect_stdout| and |redirect_stderr| to True.
302      quiet: Set |stdout_to_pipe| and |combine_stdout_stderr| to True.
303      close_fds: Whether to close all fds before running |cmd|.
304
305    Returns:
306      A CommandResult object.
307
308    Raises:
309      RunCommandError: Raises exception on error with optional error_message.
310    """
311    if capture_output:
312        redirect_stdout, redirect_stderr = True, True
313
314    if quiet:
315        stdout_to_pipe, combine_stdout_stderr = True, True
316
317    # Set default for variables.
318    stdout = None
319    stderr = None
320    stdin = None
321    cmd_result = CommandResult()
322
323    # Force the timeout to float; in the process, if it's not convertible,
324    # a self-explanatory exception will be thrown.
325    kill_timeout = float(kill_timeout)
326
327    def _get_tempfile():
328        try:
329            return tempfile.TemporaryFile(bufsize=0)
330        except EnvironmentError as e:
331            if e.errno != errno.ENOENT:
332                raise
333            # This can occur if we were pointed at a specific location for our
334            # TMP, but that location has since been deleted.  Suppress that
335            # issue in this particular case since our usage gurantees deletion,
336            # and since this is primarily triggered during hard cgroups
337            # shutdown.
338            return tempfile.TemporaryFile(bufsize=0, dir='/tmp')
339
340    # Modify defaults based on parameters.
341    # Note that tempfiles must be unbuffered else attempts to read
342    # what a separate process did to that file can result in a bad
343    # view of the file.
344    # The Popen API accepts either an int or a file handle for stdout/stderr.
345    # pylint: disable=redefined-variable-type
346    if log_stdout_to_file:
347        stdout = open(log_stdout_to_file, 'w+')
348    elif stdout_to_pipe:
349        stdout = subprocess.PIPE
350    elif redirect_stdout:
351        stdout = _get_tempfile()
352
353    if combine_stdout_stderr:
354        stderr = subprocess.STDOUT
355    elif redirect_stderr:
356        stderr = _get_tempfile()
357    # pylint: enable=redefined-variable-type
358
359    # If subprocesses have direct access to stdout or stderr, they can bypass
360    # our buffers, so we need to flush to ensure that output is not interleaved.
361    if stdout is None or stderr is None:
362        sys.stdout.flush()
363        sys.stderr.flush()
364
365    # If input is a string, we'll create a pipe and send it through that.
366    # Otherwise we assume it's a file object that can be read from directly.
367    if isinstance(input, basestring):
368        stdin = subprocess.PIPE
369    elif input is not None:
370        stdin = input
371        input = None
372
373    if isinstance(cmd, basestring):
374        if not shell:
375            raise Exception('Cannot run a string command without a shell')
376        cmd = ['/bin/bash', '-c', cmd]
377        shell = False
378    elif shell:
379        raise Exception('Cannot run an array command with a shell')
380
381    # If we are using enter_chroot we need to use enterchroot pass env through
382    # to the final command.
383    env = env.copy() if env is not None else os.environ.copy()
384    env.update(extra_env if extra_env else {})
385
386    cmd_result.cmd = cmd
387
388    proc = None
389    # Verify that the signals modules is actually usable, and won't segfault
390    # upon invocation of getsignal.  See signals.SignalModuleUsable for the
391    # details and upstream python bug.
392    use_signals = rh.signals.signal_module_usable()
393    try:
394        proc = _Popen(cmd, cwd=cwd, stdin=stdin, stdout=stdout,
395                      stderr=stderr, shell=False, env=env,
396                      close_fds=close_fds)
397
398        if use_signals:
399            old_sigint = signal.getsignal(signal.SIGINT)
400            if ignore_sigint:
401                handler = signal.SIG_IGN
402            else:
403                handler = functools.partial(
404                    _kill_child_process, proc, int_timeout, kill_timeout, cmd,
405                    old_sigint)
406            signal.signal(signal.SIGINT, handler)
407
408            old_sigterm = signal.getsignal(signal.SIGTERM)
409            handler = functools.partial(_kill_child_process, proc, int_timeout,
410                                        kill_timeout, cmd, old_sigterm)
411            signal.signal(signal.SIGTERM, handler)
412
413        try:
414            (cmd_result.output, cmd_result.error) = proc.communicate(input)
415        finally:
416            if use_signals:
417                signal.signal(signal.SIGINT, old_sigint)
418                signal.signal(signal.SIGTERM, old_sigterm)
419
420            if stdout and not log_stdout_to_file and not stdout_to_pipe:
421                # The linter is confused by how stdout is a file & an int.
422                # pylint: disable=maybe-no-member,no-member
423                stdout.seek(0)
424                cmd_result.output = stdout.read()
425                stdout.close()
426
427            if stderr and stderr != subprocess.STDOUT:
428                # The linter is confused by how stderr is a file & an int.
429                # pylint: disable=maybe-no-member,no-member
430                stderr.seek(0)
431                cmd_result.error = stderr.read()
432                stderr.close()
433
434        cmd_result.returncode = proc.returncode
435
436        if not error_code_ok and proc.returncode:
437            msg = 'cwd=%s' % cwd
438            if extra_env:
439                msg += ', extra env=%s' % extra_env
440            if error_message:
441                msg += '\n%s' % error_message
442            raise RunCommandError(msg, cmd_result)
443    except OSError as e:
444        estr = str(e)
445        if e.errno == errno.EACCES:
446            estr += '; does the program need `chmod a+x`?'
447        if error_code_ok:
448            cmd_result = CommandResult(cmd=cmd, error=estr, returncode=255)
449        else:
450            raise RunCommandError(estr, CommandResult(cmd=cmd), exception=e)
451    finally:
452        if proc is not None:
453            # Ensure the process is dead.
454            _kill_child_process(proc, int_timeout, kill_timeout, cmd, None,
455                                None, None)
456
457    return cmd_result
458# pylint: enable=redefined-builtin
459
460
461def collection(classname, **kwargs):
462    """Create a new class with mutable named members.
463
464    This is like collections.namedtuple, but mutable.  Also similar to the
465    python 3.3 types.SimpleNamespace.
466
467    Example:
468      # Declare default values for this new class.
469      Foo = collection('Foo', a=0, b=10)
470      # Create a new class but set b to 4.
471      foo = Foo(b=4)
472      # Print out a (will be the default 0) and b (will be 4).
473      print('a = %i, b = %i' % (foo.a, foo.b))
474    """
475
476    def sn_init(self, **kwargs):
477        """The new class's __init__ function."""
478        # First verify the kwargs don't have excess settings.
479        valid_keys = set(self.__slots__[1:])
480        these_keys = set(kwargs.keys())
481        invalid_keys = these_keys - valid_keys
482        if invalid_keys:
483            raise TypeError('invalid keyword arguments for this object: %r' %
484                            invalid_keys)
485
486        # Now initialize this object.
487        for k in valid_keys:
488            setattr(self, k, kwargs.get(k, self.__defaults__[k]))
489
490    def sn_repr(self):
491        """The new class's __repr__ function."""
492        return '%s(%s)' % (classname, ', '.join(
493            '%s=%r' % (k, getattr(self, k)) for k in self.__slots__[1:]))
494
495    # Give the new class a unique name and then generate the code for it.
496    classname = 'Collection_%s' % classname
497    expr = '\n'.join((
498        'class %(classname)s(object):',
499        '  __slots__ = ["__defaults__", "%(slots)s"]',
500        '  __defaults__ = {}',
501    )) % {
502        'classname': classname,
503        'slots': '", "'.join(sorted(str(k) for k in kwargs)),
504    }
505
506    # Create the class in a local namespace as exec requires.
507    namespace = {}
508    exec expr in namespace  # pylint: disable=exec-used
509    new_class = namespace[classname]
510
511    # Bind the helpers.
512    new_class.__defaults__ = kwargs.copy()
513    new_class.__init__ = sn_init
514    new_class.__repr__ = sn_repr
515
516    return new_class
517