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