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