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         uid = os.stat(os.getcwd()).st_uid
 198         # consider also checking ruid?
 199         if uid != os.geteuid():
 200             kwargs['uid'] = uid
 201             return self.callAsUser(*args, **kwargs)
 202         else:
 203             return self.call(*args, **kwargs)
 204     def eval(self, *args, **kwargs):
 205         """
 206         Evaluates a command and returns its output, with trailing newlines
 207         stripped (like backticks in Bash).  This is a convenience method for
 208         calling :meth:`call` with ``strip``.
 209
 210             >>> sh = Shell()
 211             >>> sh.eval("echo", "Foobar")
 212             'Foobar'
 213         """
 214         kwargs["strip"] = True
 215         return self.call(*args, **kwargs)
 216
 217 class ParallelShell(Shell):
 218     """
 219     Modifies the semantics of :class:`Shell` so that
 220     commands are queued here, and executed in parallel using waitpid
 221     with ``max`` subprocesses, and result in callback execution
 222     when they finish.
 223
 224     .. method:: call(*args, **kwargs)
 225
 226         Enqueues a system call for parallel processing.  If there are
 227         no openings in the queue, this will block.  Keyword arguments
 228         are the same as :meth:`Shell.call` with the following additions:
 229
 230         :param on_success: Callback function for success (zero exit status).
 231             The callback function should accept two arguments,
 232             ``stdout`` and ``stderr``.
 233         :param on_error: Callback function for failure (nonzero exit status).
 234             The callback function should accept one argument, the
 235             exception that would have been thrown by the synchronous
 236             version.
 237         :return: The :class:`subprocess.Proc` object that was opened.
 238
 239     .. method:: callAsUser(*args, **kwargs)
 240
 241         Enqueues a system call under a different user for parallel
 242         processing.  Keyword arguments are the same as
 243         :meth:`Shell.callAsUser` with the additions of keyword
 244         arguments from :meth:`call`.
 245
 246     .. method:: safeCall(*args, **kwargs)
 247
 248         Enqueues a "safe" call for parallel processing.  Keyword
 249         arguments are the same as :meth:`Shell.safeCall` with the
 250         additions of keyword arguments from :meth:`call`.
 251
 252     .. method:: eval(*args, **kwargs)
 253
 254         No difference from :meth:`call`.  Consider having a
 255         non-parallel shell if the program you are shelling out
 256         to is fast.
 257
 258     """
 259     def __init__(self, dry = False, max = 10):
 260         super(ParallelShell, self).__init__(dry=dry)
 261         self.running = {}
 262         self.max = max # maximum of commands to run in parallel
 263     @staticmethod
 264     def make(no_parallelize, max):
 265         """Convenience method oriented towards command modules."""
 266         if no_parallelize:
 267             return DummyParallelShell()
 268         else:
 269             return ParallelShell(max=max)
 270     def _async(self, proc, args, python, on_success, on_error, **kwargs):
 271         """
 272         Gets handed a :class:`subprocess.Proc` object from our deferred
 273         execution.  See :meth:`Shell.call` source code for details.
 274         """
 275         self.running[proc.pid] = (proc, args, python, on_success, on_error)
 276         return True # so that the parent function returns
 277     def _wait(self):
 278         """
 279         Blocking call that waits for an open subprocess slot.  This is
 280         automatically called by :meth:`Shell.call`.
 281         """
 282         # XXX: This API sucks; the actuall call/callAsUser call should
 283         # probably block automatically (unless I have a good reason not to)
 284         # bail out immediately on initial ramp up
 285         if len(self.running) < self.max: return
 286         # now, wait for open pids.
 287         try:
 288             self.reap(*os.waitpid(-1, 0))
 289         except OSError as e:
 290             if e.errno == errno.ECHILD: return
 291             raise
 292     def join(self):
 293         """Waits for all of our subprocesses to terminate."""
 294         try:
 295             while True:
 296                 self.reap(*os.waitpid(-1, 0))
 297         except OSError as e:
 298             if e.errno == errno.ECHILD: return
 299             raise
 300     def reap(self, pid, status):
 301         """Reaps a process."""
 302         # ooh, zombie process. reap it
 303         proc, args, python, on_success, on_error = self.running.pop(pid)
 304         # XXX: this is slightly dangerous; should actually use
 305         # temporary files
 306         stdout = proc.stdout.read()
 307         stderr = proc.stderr.read()
 308         self._log(stdout, stderr)
 309         if status:
 310             if python: eclass = PythonCallError
 311             else: eclass = CallError
 312             on_error(eclass(proc.returncode, args, stdout, stderr))
 313             return
 314         on_success(stdout, stderr)
 315
 316 # Setup a convenience global instance
 317 shell = Shell()
 318 call = shell.call
 319 callAsUser = shell.callAsUser
 320 safeCall = shell.safeCall
 321 eval = shell.eval
 322
 323 class DummyParallelShell(ParallelShell):
 324     """Same API as :class:`ParallelShell`, but doesn't actually
 325     parallelize (i.e. all calls to :meth:`wait` block.)"""
 326     def __init__(self, dry = False):
 327         super(DummyParallelShell, self).__init__(dry=dry, max=1)
 328
 329 class Error(wizard.Error):
 330     """Base exception for this module"""
 331     pass
 332
 333 class CallError(Error):
 334     """Indicates that a subprocess call returned a nonzero exit status."""
 335     #: The exit code of the failed subprocess.
 336     code = None
 337     #: List of the program and arguments that failed.
 338     args = None
 339     #: The stdout of the program.
 340     stdout = None
 341     #: The stderr of the program.
 342     stderr = None
 343     def __init__(self, code, args, stdout, stderr):
 344         self.code = code
 345         self.args = args
 346         self.stdout = stdout
 347         self.stderr = stderr
 348     def __str__(self):
 349         compact = self.stderr.rstrip().split("\n")[-1]
 350         return "%s (exited with %d)\n%s" % (compact, self.code, self.stderr)
 351
 352 class PythonCallError(CallError):
 353     """
 354     Indicates that a Python subprocess call had an uncaught exception.
 355     This exception also contains the attributes of :class:`CallError`.
 356     """
 357     #: Name of the uncaught exception.
 358     name = None
 359     def __init__(self, code, args, stdout, stderr):
 360         if stderr: self.name = util.get_exception_name(stderr)
 361         CallError.__init__(self, code, args, stdout, stderr)
 362     def __str__(self):
 363         if self.name:
 364             return "PythonCallError [%s]\n%s" % (self.name, self.stderr)
 365         else:
 366             return "PythonCallError\n%s" % self.stderr
 367