• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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
34import rh.shell
35import rh.signals
36
37
38class 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
53class 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
100class 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
108def 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
165def _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
210class _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
260def 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
457def 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