wizard/shell.py

   1 """
   2 Wrappers around subprocess functionality that simulate an actual shell.
   3
   4 .. testsetup:: *
   5
   6     from wizard.shell import *
   7 """
   8
   9 import subprocess
  10 import logging
  11 import sys
  12 import os
  13
  14 import wizard
  15 from wizard import util
  16
  17 wizard_bin = sys.argv[0]
  18 """
  19 This is the path to the wizard executable as specified
  20 by the caller; it lets us recursively invoke wizard.
  21 """
  22
  23 def is_python(args):
  24     """Detects whether or not an argument list invokes a Python program."""
  25     return args[0] == "python" or args[0] == "wizard"
  26
  27 class Shell(object):
  28     """
  29     An advanced shell that performs logging.  If ``dry`` is ``True``,
  30     no commands are actually run.
  31     """
  32     def __init__(self, dry = False):
  33         self.dry = dry
  34     def call(self, *args, **kwargs):
  35         """
  36         Performs a system call.  The actual executable and options should
  37         be passed as arguments to this function.  It will magically
  38         ensure that 'wizard' as a command works. Several keyword arguments
  39         are also supported:
  40
  41         :param python: explicitly marks the subprocess as Python or not Python
  42             for improved error reporting.  By default, we use
  43             :func:`is_python` to autodetect this.
  44         :param input: input to feed the subprocess on standard input.
  45         :param interactive: whether or not directly hook up all pipes
  46             to the controlling terminal, to allow interaction with subprocess.
  47         :returns: a tuple of strings ``(stdout, stderr)``
  48
  49         >>> sh = Shell()
  50         >>> sh.call("echo", "Foobar")
  51         ('Foobar\\n', '')
  52
  53         .. note::
  54
  55             This function does not munge trailing whitespace.  A common
  56             idiom for dealing with this is::
  57
  58                 sh.call("echo", "Foobar")[0].rstrip()
  59         """
  60         kwargs.setdefault("python", None)
  61         logging.info("Running `" + ' '.join(args) + "`")
  62         if self.dry:
  63             return
  64         if kwargs["python"] is None and is_python(args):
  65             kwargs["python"] = True
  66         if args[0] == "wizard":
  67             args = list(args)
  68             args[0] = wizard_bin
  69         kwargs.setdefault("input", None)
  70         kwargs.setdefault("interactive", False)
  71         if kwargs["interactive"]:
  72             stdout=sys.stdout
  73             stdin=sys.stdin
  74             stderr=sys.stderr
  75         else:
  76             stdout=subprocess.PIPE
  77             stdin=subprocess.PIPE
  78             stderr=subprocess.PIPE
  79         # XXX: There is a possible problem here where we can fill up
  80         # the kernel buffer if we have 64KB of data.  This shouldn't
  81         # be a problem, and the fix for such case would be to write to
  82         # temporary files instead of a pipe.
  83         # Another possible way of fixing this is converting from a
  84         # waitpid() pump to a select() pump, creating a pipe to
  85         # ourself, and then setting up a
  86         # SIGCHILD handler to write a single byte to the pipe to get
  87         # us out of select() when a subprocess exits.
  88         proc = subprocess.Popen(args, stdout=stdout, stderr=stderr, stdin=stdin)
  89         if hasattr(self, "_async"):
  90             self._async(proc, args, **kwargs)
  91             return proc
  92         stdout, stderr = proc.communicate(kwargs["input"])
  93         if not kwargs["interactive"]:
  94             self._log(stdout, stderr)
  95         if proc.returncode:
  96             if kwargs["python"]: eclass = PythonCallError
  97             else: eclass = CallError
  98             raise eclass(proc.returncode, args, stdout, stderr)
  99         return (stdout, stderr)
 100     def _log(self, stdout, stderr):
 101         """Logs the standard output and standard input from a command."""
 102         if stdout:
 103             logging.debug("STDOUT:\n" + stdout)
 104         if stderr:
 105             logging.debug("STDERR:\n" + stderr)
 106     def callAsUser(self, *args, **kwargs):
 107         """
 108         Performs a system call as a different user.  This is only possible
 109         if you are running as root.  Keyword arguments
 110         are the same as :meth:`call` with the following additions:
 111
 112         :param user: name of the user to run command as.
 113         :param uid: uid of the user to run command as.
 114
 115         .. note::
 116
 117             The resulting system call internally uses :command:`sudo`,
 118             and as such environment variables will get scrubbed.  We
 119             manually preserve :envvar:`SSH_GSSAPI_NAME`.
 120         """
 121         user = kwargs.pop("user", None)
 122         uid = kwargs.pop("uid", None)
 123         kwargs.setdefault("python", is_python(args))
 124         if not user and not uid: return self.call(*args, **kwargs)
 125         if util.get_operator_name():
 126             # This might be generalized as "preserve some environment"
 127             args.insert(0, "SSH_GSSAPI_NAME=" + util.get_operator_name())
 128         if uid: return self.call("sudo", "-u", "#" + str(uid), *args, **kwargs)
 129         if user: return self.call("sudo", "-u", user, *args, **kwargs)
 130     def safeCall(self, *args, **kwargs):
 131         """
 132         Checks if the owner of the current working directory is the same
 133         as the current user, and if it isn't, attempts to sudo to be
 134         that user.  The intended use case is for calling Git commands
 135         when running as root, but this method should be used when
 136         interfacing with any moderately complex program that depends
 137         on working directory context.  Keyword arguments are the
 138         same as :meth:`call`.
 139         """
 140         uid = os.stat(os.getcwd()).st_uid
 141         # consider also checking ruid?
 142         if uid != os.geteuid():
 143             kwargs['uid'] = uid
 144             return self.callAsUser(*args, **kwargs)
 145         else:
 146             return self.call(*args, **kwargs)
 147
 148 class ParallelShell(Shell):
 149     """
 150     Modifies the semantics of :class:`Shell` so that
 151     commands are queued here, and executed in parallel using waitpid
 152     with ``max`` subprocesses, and result in callback execution
 153     when they finish.
 154
 155     Before enqueueing a system call with :meth:`call` or :meth:`callAsUser`,
 156     you should wait for an open slot using :meth:`wait`; otherwise,
 157     ``max`` rate limiting will have no effect.  For example::
 158
 159         sh = ParallelShell()
 160         for command in commands_to_execute_in_parallel:
 161             sh.wait()
 162             sh.call(*command)
 163         sh.join()
 164
 165     .. method:: call(*args, **kwargs)
 166
 167         Enqueues a system call for parallel processing.  Keyword arguments
 168         are the same as :meth:`Shell.call` with the following additions:
 169
 170         :param on_success: Callback function for success (zero exit status).
 171             The callback function should accept two arguments,
 172             ``stdout`` and ``stderr``.
 173         :param on_error: Callback function for failure (nonzero exit status).
 174             The callback function should accept one argument, the
 175             exception that would have been thrown by the synchronous
 176             version.
 177         :return: The :class:`subprocess.Proc` object that was opened.
 178
 179     .. method:: callAsUser(*args, **kwargs)
 180
 181         Enqueues a system call under a different user for parallel
 182         processing.  Keyword arguments are the same as
 183         :meth:`Shell.callAsUser` with the additions of keyword
 184         arguments from :meth:`call`.
 185     """
 186     def __init__(self, dry = False, max = 10):
 187         super(ParallelShell, self).__init__(dry=dry)
 188         self.running = {}
 189         self.max = max # maximum of commands to run in parallel
 190     def _async(self, proc, args, python, on_success, on_error):
 191         """
 192         Gets handed a :class:`subprocess.Proc` object from our deferred
 193         execution.  See :meth:`Shell.call` source code for details.
 194         """
 195         self.running[proc.pid] = (proc, args, python, on_success, on_error)
 196     def wait(self):
 197         """
 198         Blocking call that waits for an open subprocess slot.  You should
 199         call this before enqueuing.
 200
 201         .. note::
 202
 203             This method may become unnecessary in the future.
 204         """
 205         # XXX: This API sucks; the actuall call/callAsUser call should
 206         # probably block automatically (unless I have a good reason not to)
 207         # bail out immediately on initial ramp up
 208         if len(self.running) < self.max: return
 209         # now, wait for open pids.
 210         try:
 211             pid, status = os.waitpid(-1, 0)
 212         except OSError as e:
 213             if e.errno == errno.ECHILD: return
 214             raise e
 215         # ooh, zombie process. reap it
 216         proc, args, python, on_success, on_error = self.running.pop(pid)
 217         # XXX: this is slightly dangerous; should actually use
 218         # temporary files
 219         stdout = proc.stdout.read()
 220         stderr = proc.stderr.read()
 221         self._log(stdout, stderr)
 222         if status:
 223             if python: eclass = PythonCallError
 224             else: eclass = CallError
 225             on_error(eclass(proc.returncode, args, stdout, stderr))
 226             return
 227         on_success(stdout, stderr)
 228     def join(self):
 229         """Waits for all of our subprocesses to terminate."""
 230         try:
 231             while os.waitpid(-1, 0):
 232                 pass
 233         except OSError as e:
 234             if e.errno == errno.ECHILD: return
 235             raise e
 236
 237 class DummyParallelShell(ParallelShell):
 238     """Same API as :class:`ParallelShell`, but doesn't actually
 239     parallelize (i.e. all calls to :meth:`wait` block.)"""
 240     def __init__(self, dry = False):
 241         super(DummyParallelShell, self).__init__(dry=dry, max=1)
 242
 243 class Error(wizard.Error):
 244     """Base exception for this module"""
 245     pass
 246
 247 class CallError(Error):
 248     """Indicates that a subprocess call returned a nonzero exit status."""
 249     #: The exit code of the failed subprocess.
 250     code = None
 251     #: List of the program and arguments that failed.
 252     args = None
 253     #: The stdout of the program.
 254     stdout = None
 255     #: The stderr of the program.
 256     stderr = None
 257     def __init__(self, code, args, stdout, stderr):
 258         self.code = code
 259         self.args = args
 260         self.stdout = stdout
 261         self.stderr = stderr
 262     def __str__(self):
 263         return "CallError [%d]" % self.code
 264
 265 class PythonCallError(CallError):
 266     """
 267     Indicates that a Python subprocess call had an uncaught exception.
 268     This exception also contains the attributes of :class:`CallError`.
 269     """
 270     #: Name of the uncaught exception.
 271     name = None
 272     def __init__(self, code, args, stdout, stderr):
 273         if stderr: self.name = util.get_exception_name(stderr)
 274         CallError.__init__(self, code, args, stdout, stderr)
 275     def __str__(self):
 276         if self.name:
 277             return "PythonCallError [%s]" % self.name
 278         else:
 279             return "PythonCallError"
 280