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 import errno
  14
  15 import wizard
  16 from wizard import util
  17
  18 wizard_bin = sys.argv[0]
  19 """
  20 This is the path to the wizard executable as specified
  21 by the caller; it lets us recursively invoke wizard.
  22 """
  23
  24 def is_python(args):
  25     """Detects whether or not an argument list invokes a Python program."""
  26     return args[0] == "python" or args[0] == "wizard"
  27
  28 def drop_priviledges(dir, log_file):
  29     """
  30     Checks if we are running as root.  If we are, attempt to drop
  31     priviledges to the user who owns ``dir``, by re-calling
  32     itself using sudo with exec, such that the new process subsumes our
  33     current one.  If ``log_file`` is passed, the file is chown'ed
  34     to the user we are dropping priviledges to, so the subprocess
  35     can write to it.
  36     """
  37     if os.getuid():
  38         return
  39     uid = util.get_dir_uid(dir)
  40     if not uid:
  41         return
  42     args = []
  43     for k,v in os.environ.items():
  44         if k.startswith('WIZARD_') or k == "SSH_GSSAPI_NAME":
  45             args.append("%s=%s" % (k,v))
  46     args += sys.argv
  47     logging.debug("Dropping priviledges")
  48     if log_file: os.chown(log_file, uid, -1)
  49     os.execlp('sudo', 'sudo', '-u', '#' + str(uid), *args)
  50
  51 class Shell(object):
  52     """
  53     An advanced shell that performs logging.  If ``dry`` is ``True``,
  54     no commands are actually run.
  55     """
  56     def __init__(self, dry = False):
  57         self.dry = dry
  58         self.cwd = None
  59     def call(self, *args, **kwargs):
  60         """
  61         Performs a system call.  The actual executable and options should
  62         be passed as arguments to this function.  It will magically
  63         ensure that 'wizard' as a command works. Several keyword arguments
  64         are also supported:
  65
  66         :param python: explicitly marks the subprocess as Python or not Python
  67             for improved error reporting.  By default, we use
  68             :func:`is_python` to autodetect this.
  69         :param input: input to feed the subprocess on standard input.
  70         :param interactive: whether or not directly hook up all pipes
  71             to the controlling terminal, to allow interaction with subprocess.
  72         :param strip: if ``True``, instead of returning a tuple,
  73             return the string stdout output of the command with trailing newlines
  74             removed.  This emulates the behavior of backticks and ``$()`` in Bash.
  75             Prefer to use :meth:`eval` instead (you should only need to explicitly
  76             specify this if you are using another wrapper around this function).
  77         :param log: if True, we log the call as INFO, if False, we log the call
  78             as DEBUG, otherwise, we detect based on ``strip``.
  79         :param stdout:
  80         :param stderr:
  81         :param stdin: a file-type object that will be written to or read from as a pipe.
  82         :returns: a tuple of strings ``(stdout, stderr)``, or a string ``stdout``
  83             if ``strip`` is specified.
  84
  85         >>> sh = Shell()
  86         >>> sh.call("echo", "Foobar")
  87         ('Foobar\\n', '')
  88         >>> sh.call("cat", input='Foobar')
  89         ('Foobar', '')
  90         """
  91         self._wait()
  92         kwargs.setdefault("interactive", False)
  93         kwargs.setdefault("strip", False)
  94         kwargs.setdefault("python", None)
  95         kwargs.setdefault("log", None)
  96         kwargs.setdefault("stdout", subprocess.PIPE)
  97         kwargs.setdefault("stdin", subprocess.PIPE)
  98         kwargs.setdefault("stderr", subprocess.PIPE)
  99         msg = "Running `" + ' '.join(args) + "`"
 100         if kwargs["strip"] and not kwargs["log"] is True or kwargs["log"] is False:
 101             logging.debug(msg)
 102         else:
 103             logging.info(msg)
 104         if self.dry:
 105             if kwargs["strip"]:
 106                 return ''
 107             return None, None
 108         if kwargs["python"] is None and is_python(args):
 109             kwargs["python"] = True
 110         if args[0] == "wizard":
 111             args = list(args)
 112             args[0] = wizard_bin
 113         kwargs.setdefault("input", None)
 114         if kwargs["interactive"]:
 115             stdout=sys.stdout
 116             stdin=sys.stdin
 117             stderr=sys.stderr
 118         else:
 119             stdout=kwargs["stdout"]
 120             stdin=kwargs["stdin"]
 121             stderr=kwargs["stderr"]
 122         # XXX: There is a possible problem here where we can fill up
 123         # the kernel buffer if we have 64KB of data.  This shouldn't
 124         # be a problem, and the fix for such case would be to write to
 125         # temporary files instead of a pipe.
 126         # Another possible way of fixing this is converting from a
 127         # waitpid() pump to a select() pump, creating a pipe to
 128         # ourself, and then setting up a
 129         # SIGCHILD handler to write a single byte to the pipe to get
 130         # us out of select() when a subprocess exits.
 131         proc = subprocess.Popen(args, stdout=stdout, stderr=stderr, stdin=stdin, cwd=self.cwd, )
 132         if self._async(proc, args, **kwargs):
 133             return proc
 134         stdout, stderr = proc.communicate(kwargs["input"])
 135         # can occur if we were doing interactive communication; i.e.
 136         # we didn't pass in PIPE.
 137         if stdout is None:
 138             stdout = ""
 139         if stderr is None:
 140             stderr = ""
 141         if not kwargs["interactive"]:
 142             if kwargs["strip"]:
 143                 self._log(None, stderr)
 144             else:
 145                 self._log(stdout, stderr)
 146         if proc.returncode:
 147             if kwargs["python"]: eclass = PythonCallError
 148             else: eclass = CallError
 149             raise eclass(proc.returncode, args, stdout, stderr)
 150         if kwargs["strip"]:
 151             return str(stdout).rstrip("\n")
 152         return (stdout, stderr)
 153     def _log(self, stdout, stderr):
 154         """Logs the standard output and standard input from a command."""
 155         if stdout:
 156             logging.debug("STDOUT:\n" + stdout)
 157         if stderr:
 158             logging.debug("STDERR:\n" + stderr)
 159     def _wait(self):
 160         pass
 161     def _async(self, *args, **kwargs):
 162         return False
 163     def callAsUser(self, *args, **kwargs):
 164         """
 165         Performs a system call as a different user.  This is only possible
 166         if you are running as root.  Keyword arguments
 167         are the same as :meth:`call` with the following additions:
 168
 169         :param user: name of the user to run command as.
 170         :param uid: uid of the user to run command as.
 171
 172         .. note::
 173
 174             The resulting system call internally uses :command:`sudo`,
 175             and as such environment variables will get scrubbed.  We
 176             manually preserve :envvar:`SSH_GSSAPI_NAME`.
 177         """
 178         user = kwargs.pop("user", None)
 179         uid = kwargs.pop("uid", None)
 180         kwargs.setdefault("python", is_python(args))
 181         if not user and not uid: return self.call(*args, **kwargs)
 182         if os.getenv("SSH_GSSAPI_NAME"):
 183             # This might be generalized as "preserve some environment"
 184             args = list(args)
 185             args.insert(0, "SSH_GSSAPI_NAME=" + os.getenv("SSH_GSSAPI_NAME"))
 186         if uid: return self.call("sudo", "-u", "#" + str(uid), *args, **kwargs)
 187         if user: return self.call("sudo", "-u", user, *args, **kwargs)
 188     def safeCall(self, *args, **kwargs):
 189         """
 190         Checks if the owner of the current working directory is the same
 191         as the current user, and if it isn't, attempts to sudo to be
 192         that user.  The intended use case is for calling Git commands
 193         when running as root, but this method should be used when
 194         interfacing with any moderately complex program that depends
 195         on working directory context.  Keyword arguments are the
 196         same as :meth:`call`.
 197         """
 198         if os.getuid():
 199             return self.call(*args, **kwargs)
 200         uid = os.stat(os.getcwd()).st_uid
 201         # consider also checking ruid?
 202         if uid != os.geteuid():
 203             kwargs['uid'] = uid
 204             return self.callAsUser(*args, **kwargs)
 205         else:
 206             return self.call(*args, **kwargs)
 207     def eval(self, *args, **kwargs):
 208         """
 209         Evaluates a command and returns its output, with trailing newlines
 210         stripped (like backticks in Bash).  This is a convenience method for
 211         calling :meth:`call` with ``strip``.
 212
 213             >>> sh = Shell()
 214             >>> sh.eval("echo", "Foobar")
 215             'Foobar'
 216         """
 217         kwargs["strip"] = True
 218         return self.call(*args, **kwargs)
 219     def setcwd(self, cwd):
 220         """
 221         Sets the directory processes are executed in. This sets a value
 222         to be passed as the ``cwd`` argument to ``subprocess.Popen``.
 223         """
 224         self.cwd = cwd
 225
 226 class ParallelShell(Shell):
 227     """
 228     Modifies the semantics of :class:`Shell` so that
 229     commands are queued here, and executed in parallel using waitpid
 230     with ``max`` subprocesses, and result in callback execution
 231     when they finish.
 232
 233     .. method:: call(*args, **kwargs)
 234
 235         Enqueues a system call for parallel processing.  If there are
 236         no openings in the queue, this will block.  Keyword arguments
 237         are the same as :meth:`Shell.call` with the following additions:
 238
 239         :param on_success: Callback function for success (zero exit status).
 240             The callback function should accept two arguments,
 241             ``stdout`` and ``stderr``.
 242         :param on_error: Callback function for failure (nonzero exit status).
 243             The callback function should accept one argument, the
 244             exception that would have been thrown by the synchronous
 245             version.
 246         :return: The :class:`subprocess.Proc` object that was opened.
 247
 248     .. method:: callAsUser(*args, **kwargs)
 249
 250         Enqueues a system call under a different user for parallel
 251         processing.  Keyword arguments are the same as
 252         :meth:`Shell.callAsUser` with the additions of keyword
 253         arguments from :meth:`call`.
 254
 255     .. method:: safeCall(*args, **kwargs)
 256
 257         Enqueues a "safe" call for parallel processing.  Keyword
 258         arguments are the same as :meth:`Shell.safeCall` with the
 259         additions of keyword arguments from :meth:`call`.
 260
 261     .. method:: eval(*args, **kwargs)
 262
 263         No difference from :meth:`call`.  Consider having a
 264         non-parallel shell if the program you are shelling out
 265         to is fast.
 266
 267     """
 268     def __init__(self, dry = False, max = 10):
 269         super(ParallelShell, self).__init__(dry=dry)
 270         self.running = {}
 271         self.max = max # maximum of commands to run in parallel
 272     @staticmethod
 273     def make(no_parallelize, max):
 274         """Convenience method oriented towards command modules."""
 275         if no_parallelize:
 276             return DummyParallelShell()
 277         else:
 278             return ParallelShell(max=max)
 279     def _async(self, proc, args, python, on_success, on_error, **kwargs):
 280         """
 281         Gets handed a :class:`subprocess.Proc` object from our deferred
 282         execution.  See :meth:`Shell.call` source code for details.
 283         """
 284         self.running[proc.pid] = (proc, args, python, on_success, on_error)
 285         return True # so that the parent function returns
 286     def _wait(self):
 287         """
 288         Blocking call that waits for an open subprocess slot.  This is
 289         automatically called by :meth:`Shell.call`.
 290         """
 291         # XXX: This API sucks; the actual call/callAsUser call should
 292         # probably block automatically (unless I have a good reason not to)
 293         # bail out immediately on initial ramp up
 294         if len(self.running) < self.max: return
 295         # now, wait for open pids.
 296         try:
 297             self.reap(*os.waitpid(-1, 0))
 298         except OSError as e:
 299             if e.errno == errno.ECHILD: return
 300             raise
 301     def join(self):
 302         """Waits for all of our subprocesses to terminate."""
 303         try:
 304             while True:
 305                 self.reap(*os.waitpid(-1, 0))
 306         except OSError as e:
 307             if e.errno == errno.ECHILD: return
 308             raise
 309     def reap(self, pid, status):
 310         """Reaps a process."""
 311         # ooh, zombie process. reap it
 312         proc, args, python, on_success, on_error = self.running.pop(pid)
 313         # XXX: this is slightly dangerous; should actually use
 314         # temporary files
 315         stdout = proc.stdout.read()
 316         stderr = proc.stderr.read()
 317         self._log(stdout, stderr)
 318         if status:
 319             if python: eclass = PythonCallError
 320             else: eclass = CallError
 321             on_error(eclass(proc.returncode, args, stdout, stderr))
 322             return
 323         on_success(stdout, stderr)
 324
 325 # Setup a convenience global instance
 326 shell = Shell()
 327 call = shell.call
 328 callAsUser = shell.callAsUser
 329 safeCall = shell.safeCall
 330 eval = shell.eval
 331
 332 class DummyParallelShell(ParallelShell):
 333     """Same API as :class:`ParallelShell`, but doesn't actually
 334     parallelize (i.e. all calls to :meth:`wait` block.)"""
 335     def __init__(self, dry = False):
 336         super(DummyParallelShell, self).__init__(dry=dry, max=1)
 337
 338 class Error(wizard.Error):
 339     """Base exception for this module"""
 340     pass
 341
 342 class CallError(Error):
 343     """Indicates that a subprocess call returned a nonzero exit status."""
 344     #: The exit code of the failed subprocess.
 345     code = None
 346     #: List of the program and arguments that failed.
 347     args = None
 348     #: The stdout of the program.
 349     stdout = None
 350     #: The stderr of the program.
 351     stderr = None
 352     def __init__(self, code, args, stdout, stderr):
 353         self.code = code
 354         self.args = args
 355         self.stdout = stdout
 356         self.stderr = stderr
 357     def __str__(self):
 358         compact = self.stderr.rstrip().split("\n")[-1]
 359         return "%s (exited with %d)\n%s" % (compact, self.code, self.stderr)
 360
 361 class PythonCallError(CallError):
 362     """
 363     Indicates that a Python subprocess call had an uncaught exception.
 364     This exception also contains the attributes of :class:`CallError`.
 365     """
 366     #: Name of the uncaught exception.
 367     name = None
 368     def __init__(self, code, args, stdout, stderr):
 369         if stderr: self.name = util.get_exception_name(stderr)
 370         CallError.__init__(self, code, args, stdout, stderr)
 371     def __str__(self):
 372         if self.name:
 373             return "PythonCallError [%s]\n%s" % (self.name, self.stderr)
 374         else:
 375             return "PythonCallError\n%s" % self.stderr
 376