source: branches/fc15-dev/host/credit-card/shell.py @ 2278

"""
Wrappers around subprocess functionality that simulate an actual shell.
"""

import subprocess
import logging
import sys
import os
import errno

class Shell(object):
    """
    An advanced shell that performs logging.  If ``dry`` is ``True``,
    no commands are actually run.
    """
    def __init__(self, dry = False):
        self.dry = dry
        self.cwd = None
    def call(self, *args, **kwargs):
        """
        Performs a system call.  The actual executable and options should
        be passed as arguments to this function.  Several keyword arguments
        are also supported:

        :param input: input to feed the subprocess on standard input.
        :param interactive: whether or not directly hook up all pipes
            to the controlling terminal, to allow interaction with subprocess.
        :param strip: if ``True``, instead of returning a tuple,
            return the string stdout output of the command with trailing newlines
            removed.  This emulates the behavior of backticks and ``$()`` in Bash.
            Prefer to use :meth:`eval` instead (you should only need to explicitly
            specify this if you are using another wrapper around this function).
        :param log: if True, we log the call as INFO, if False, we log the call
            as DEBUG, otherwise, we detect based on ``strip``.
        :param stdout:
        :param stderr:
        :param stdin: a file-type object that will be written to or read from as a pipe.
        :returns: a tuple of strings ``(stdout, stderr)``, or a string ``stdout``
            if ``strip`` is specified.

        >>> sh = Shell()
        >>> sh.call("echo", "Foobar")
        ('Foobar\\n', '')
        >>> sh.call("cat", input='Foobar')
        ('Foobar', '')
        """
        self._wait()
        kwargs.setdefault("interactive", False)
        kwargs.setdefault("strip", False)
        kwargs.setdefault("python", None)
        kwargs.setdefault("log", None)
        kwargs.setdefault("stdout", subprocess.PIPE)
        kwargs.setdefault("stdin", subprocess.PIPE)
        kwargs.setdefault("stderr", subprocess.PIPE)
        msg = "Running `" + ' '.join(args) + "`"
        if kwargs["strip"] and not kwargs["log"] is True or kwargs["log"] is False:
            logging.debug(msg)
        else:
            logging.info(msg)
        if self.dry:
            if kwargs["strip"]:
                return ''
            return None, None
        kwargs.setdefault("input", None)
        if kwargs["interactive"]:
            stdout=sys.stdout
            stdin=sys.stdin
            stderr=sys.stderr
        else:
            stdout=kwargs["stdout"]
            stdin=kwargs["stdin"]
            stderr=kwargs["stderr"]
        # XXX: There is a possible problem here where we can fill up
        # the kernel buffer if we have 64KB of data.  This shouldn't
        # be a problem, and the fix for such case would be to write to
        # temporary files instead of a pipe.
        # Another possible way of fixing this is converting from a
        # waitpid() pump to a select() pump, creating a pipe to
        # ourself, and then setting up a
        # SIGCHILD handler to write a single byte to the pipe to get
        # us out of select() when a subprocess exits.
        proc = subprocess.Popen(args, stdout=stdout, stderr=stderr, stdin=stdin, cwd=self.cwd, )
        if self._async(proc, args, **kwargs):
            return proc
        stdout, stderr = proc.communicate(kwargs["input"])
        # can occur if we were doing interactive communication; i.e.
        # we didn't pass in PIPE.
        if stdout is None:
            stdout = ""
        if stderr is None:
            stderr = ""
        if not kwargs["interactive"]:
            if kwargs["strip"]:
                self._log(None, stderr)
            else:
                self._log(stdout, stderr)
        if proc.returncode:
            raise CallError(proc.returncode, args, stdout, stderr)
        if kwargs["strip"]:
            return str(stdout).rstrip("\n")
        return (stdout, stderr)
    def _log(self, stdout, stderr):
        """Logs the standard output and standard input from a command."""
        if stdout:
            logging.debug("STDOUT:\n" + stdout)
        if stderr:
            logging.debug("STDERR:\n" + stderr)
    def _wait(self):
        pass
    def _async(self, *args, **kwargs):
        return False
    def callAsUser(self, *args, **kwargs):
        """
        Performs a system call as a different user.  This is only possible
        if you are running as root.  Keyword arguments
        are the same as :meth:`call` with the following additions:

        :param user: name of the user to run command as.
        :param uid: uid of the user to run command as.

        .. note::

            The resulting system call internally uses :command:`sudo`,
            and as such environment variables will get scrubbed.  We
            manually preserve :envvar:`SSH_GSSAPI_NAME`.
        """
        user = kwargs.pop("user", None)
        uid = kwargs.pop("uid", None)
        if not user and not uid: return self.call(*args, **kwargs)
        if os.getenv("SSH_GSSAPI_NAME"):
            # This might be generalized as "preserve some environment"
            args = list(args)
            args.insert(0, "SSH_GSSAPI_NAME=" + os.getenv("SSH_GSSAPI_NAME"))
        if uid: return self.call("sudo", "-u", "#" + str(uid), *args, **kwargs)
        if user: return self.call("sudo", "-u", user, *args, **kwargs)
    def safeCall(self, *args, **kwargs):
        """
        Checks if the owner of the current working directory is the same
        as the current user, and if it isn't, attempts to sudo to be
        that user.  The intended use case is for calling Git commands
        when running as root, but this method should be used when
        interfacing with any moderately complex program that depends
        on working directory context.  Keyword arguments are the
        same as :meth:`call`.
        """
        if os.getuid():
            return self.call(*args, **kwargs)
        uid = os.stat(os.getcwd()).st_uid
        # consider also checking ruid?
        if uid != os.geteuid():
            kwargs['uid'] = uid
            return self.callAsUser(*args, **kwargs)
        else:
            return self.call(*args, **kwargs)
    def eval(self, *args, **kwargs):
        """
        Evaluates a command and returns its output, with trailing newlines
        stripped (like backticks in Bash).  This is a convenience method for
        calling :meth:`call` with ``strip``.

            >>> sh = Shell()
            >>> sh.eval("echo", "Foobar") 
            'Foobar'
        """
        kwargs["strip"] = True
        return self.call(*args, **kwargs)
    def setcwd(self, cwd):
        """
        Sets the directory processes are executed in. This sets a value
        to be passed as the ``cwd`` argument to ``subprocess.Popen``.
        """
        self.cwd = cwd

class ParallelShell(Shell):
    """
    Modifies the semantics of :class:`Shell` so that
    commands are queued here, and executed in parallel using waitpid
    with ``max`` subprocesses, and result in callback execution
    when they finish.

    .. method:: call(*args, **kwargs)

        Enqueues a system call for parallel processing.  If there are
        no openings in the queue, this will block.  Keyword arguments
        are the same as :meth:`Shell.call` with the following additions:

        :param on_success: Callback function for success (zero exit status).
            The callback function should accept two arguments,
            ``stdout`` and ``stderr``.
        :param on_error: Callback function for failure (nonzero exit status).
            The callback function should accept one argument, the
            exception that would have been thrown by the synchronous
            version.
        :return: The :class:`subprocess.Proc` object that was opened.

    .. method:: callAsUser(*args, **kwargs)

        Enqueues a system call under a different user for parallel
        processing.  Keyword arguments are the same as
        :meth:`Shell.callAsUser` with the additions of keyword
        arguments from :meth:`call`.

    .. method:: safeCall(*args, **kwargs)

        Enqueues a "safe" call for parallel processing.  Keyword
        arguments are the same as :meth:`Shell.safeCall` with the
        additions of keyword arguments from :meth:`call`.

    .. method:: eval(*args, **kwargs)

        No difference from :meth:`call`.  Consider having a
        non-parallel shell if the program you are shelling out
        to is fast.

    """
    def __init__(self, dry = False, max = 10):
        super(ParallelShell, self).__init__(dry=dry)
        self.running = {}
        self.max = max # maximum of commands to run in parallel
    @staticmethod
    def make(no_parallelize, max):
        """Convenience method oriented towards command modules."""
        if no_parallelize:
            return DummyParallelShell()
        else:
            return ParallelShell(max=max)
    def _async(self, proc, args, python, on_success, on_error, **kwargs):
        """
        Gets handed a :class:`subprocess.Proc` object from our deferred
        execution.  See :meth:`Shell.call` source code for details.
        """
        self.running[proc.pid] = (proc, args, python, on_success, on_error)
        return True # so that the parent function returns
    def _wait(self):
        """
        Blocking call that waits for an open subprocess slot.  This is
        automatically called by :meth:`Shell.call`.
        """
        # XXX: This API sucks; the actual call/callAsUser call should
        # probably block automatically (unless I have a good reason not to)
        # bail out immediately on initial ramp up
        if len(self.running) < self.max: return
        # now, wait for open pids.
        try:
            self.reap(*os.waitpid(-1, 0))
        except OSError as e:
            if e.errno == errno.ECHILD: return
            raise
    def join(self):
        """Waits for all of our subprocesses to terminate."""
        try:
            while True:
                self.reap(*os.waitpid(-1, 0))
        except OSError as e:
            if e.errno == errno.ECHILD: return
            raise
    def reap(self, pid, status):
        """Reaps a process."""
        # ooh, zombie process. reap it
        proc, args, python, on_success, on_error = self.running.pop(pid)
        # XXX: this is slightly dangerous; should actually use
        # temporary files
        stdout = proc.stdout.read()
        stderr = proc.stderr.read()
        self._log(stdout, stderr)
        if status:
            on_error(CallError(proc.returncode, args, stdout, stderr))
            return
        on_success(stdout, stderr)

# Setup a convenience global instance
shell = Shell()
call = shell.call
callAsUser = shell.callAsUser
safeCall = shell.safeCall
eval = shell.eval

class DummyParallelShell(ParallelShell):
    """Same API as :class:`ParallelShell`, but doesn't actually
    parallelize (i.e. all calls to :meth:`wait` block.)"""
    def __init__(self, dry = False):
        super(DummyParallelShell, self).__init__(dry=dry, max=1)

class CallError:
    """Indicates that a subprocess call returned a nonzero exit status."""
    #: The exit code of the failed subprocess.
    code = None
    #: List of the program and arguments that failed.
    args = None
    #: The stdout of the program.
    stdout = None
    #: The stderr of the program.
    stderr = None
    def __init__(self, code, args, stdout, stderr):
        self.code = code
        self.args = args
        self.stdout = stdout
        self.stderr = stderr
    def __str__(self):
        compact = self.stderr.rstrip().split("\n")[-1]
        return "%s (exited with %d)\n%s" % (compact, self.code, self.stderr)