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         if not kwargs["interactive"]:
 135             if kwargs["strip"]:
 136                 self._log(None, stderr)
 137             else:
 138                 self._log(stdout, stderr)
 139         if proc.returncode:
 140             if kwargs["python"]: eclass = PythonCallError
 141             else: eclass = CallError
 142             raise eclass(proc.returncode, args, stdout, stderr)
 143         if kwargs["strip"]:
 144             return str(stdout).rstrip("\n")
 145         return (stdout, stderr)
 146     def _log(self, stdout, stderr):
 147         """Logs the standard output and standard input from a command."""
 148         if stdout:
 149             logging.debug("STDOUT:\n" + stdout)
 150         if stderr:
 151             logging.debug("STDERR:\n" + stderr)
 152     def _wait(self):
 153         pass
 154     def _async(self, *args, **kwargs):
 155         return False
 156     def callAsUser(self, *args, **kwargs):
 157         """
 158         Performs a system call as a different user.  This is only possible
 159         if you are running as root.  Keyword arguments
 160         are the same as :meth:`call` with the following additions:
 161
 162         :param user: name of the user to run command as.
 163         :param uid: uid of the user to run command as.
 164
 165         .. note::
 166
 167             The resulting system call internally uses :command:`sudo`,
 168             and as such environment variables will get scrubbed.  We
 169             manually preserve :envvar:`SSH_GSSAPI_NAME`.
 170         """
 171         user = kwargs.pop("user", None)
 172         uid = kwargs.pop("uid", None)
 173         kwargs.setdefault("python", is_python(args))
 174         if not user and not uid: return self.call(*args, **kwargs)
 175         if os.getenv("SSH_GSSAPI_NAME"):
 176             # This might be generalized as "preserve some environment"
 177             args = list(args)
 178             args.insert(0, "SSH_GSSAPI_NAME=" + os.getenv("SSH_GSSAPI_NAME"))
 179         if uid: return self.call("sudo", "-u", "#" + str(uid), *args, **kwargs)
 180         if user: return self.call("sudo", "-u", user, *args, **kwargs)
 181     def safeCall(self, *args, **kwargs):
 182         """
 183         Checks if the owner of the current working directory is the same
 184         as the current user, and if it isn't, attempts to sudo to be
 185         that user.  The intended use case is for calling Git commands
 186         when running as root, but this method should be used when
 187         interfacing with any moderately complex program that depends
 188         on working directory context.  Keyword arguments are the
 189         same as :meth:`call`.
 190         """
 191         uid = os.stat(os.getcwd()).st_uid
 192         # consider also checking ruid?
 193         if uid != os.geteuid():
 194             kwargs['uid'] = uid
 195             return self.callAsUser(*args, **kwargs)
 196         else:
 197             return self.call(*args, **kwargs)
 198     def eval(self, *args, **kwargs):
 199         """
 200         Evaluates a command and returns its output, with trailing newlines
 201         stripped (like backticks in Bash).  This is a convenience method for
 202         calling :meth:`call` with ``strip``.
 203
 204             >>> sh = Shell()
 205             >>> sh.eval("echo", "Foobar")
 206             'Foobar'
 207         """
 208         kwargs["strip"] = True
 209         return self.call(*args, **kwargs)
 210
 211 class ParallelShell(Shell):
 212     """
 213     Modifies the semantics of :class:`Shell` so that
 214     commands are queued here, and executed in parallel using waitpid
 215     with ``max`` subprocesses, and result in callback execution
 216     when they finish.
 217
 218     .. method:: call(*args, **kwargs)
 219
 220         Enqueues a system call for parallel processing.  If there are
 221         no openings in the queue, this will block.  Keyword arguments
 222         are the same as :meth:`Shell.call` with the following additions:
 223
 224         :param on_success: Callback function for success (zero exit status).
 225             The callback function should accept two arguments,
 226             ``stdout`` and ``stderr``.
 227         :param on_error: Callback function for failure (nonzero exit status).
 228             The callback function should accept one argument, the
 229             exception that would have been thrown by the synchronous
 230             version.
 231         :return: The :class:`subprocess.Proc` object that was opened.
 232
 233     .. method:: callAsUser(*args, **kwargs)
 234
 235         Enqueues a system call under a different user for parallel
 236         processing.  Keyword arguments are the same as
 237         :meth:`Shell.callAsUser` with the additions of keyword
 238         arguments from :meth:`call`.
 239
 240     .. method:: safeCall(*args, **kwargs)
 241
 242         Enqueues a "safe" call for parallel processing.  Keyword
 243         arguments are the same as :meth:`Shell.safeCall` with the
 244         additions of keyword arguments from :meth:`call`.
 245
 246     .. method:: eval(*args, **kwargs)
 247
 248         No difference from :meth:`call`.  Consider having a
 249         non-parallel shell if the program you are shelling out
 250         to is fast.
 251
 252     """
 253     def __init__(self, dry = False, max = 10):
 254         super(ParallelShell, self).__init__(dry=dry)
 255         self.running = {}
 256         self.max = max # maximum of commands to run in parallel
 257     @staticmethod
 258     def make(no_parallelize, max):
 259         """Convenience method oriented towards command modules."""
 260         if no_parallelize:
 261             return DummyParallelShell()
 262         else:
 263             return ParallelShell(max=max)
 264     def _async(self, proc, args, python, on_success, on_error, **kwargs):
 265         """
 266         Gets handed a :class:`subprocess.Proc` object from our deferred
 267         execution.  See :meth:`Shell.call` source code for details.
 268         """
 269         self.running[proc.pid] = (proc, args, python, on_success, on_error)
 270         return True # so that the parent function returns
 271     def _wait(self):
 272         """
 273         Blocking call that waits for an open subprocess slot.  This is
 274         automatically called by :meth:`Shell.call`.
 275         """
 276         # XXX: This API sucks; the actuall call/callAsUser call should
 277         # probably block automatically (unless I have a good reason not to)
 278         # bail out immediately on initial ramp up
 279         if len(self.running) < self.max: return
 280         # now, wait for open pids.
 281         try:
 282             self.reap(*os.waitpid(-1, 0))
 283         except OSError as e:
 284             if e.errno == errno.ECHILD: return
 285             raise
 286     def join(self):
 287         """Waits for all of our subprocesses to terminate."""
 288         try:
 289             while True:
 290                 self.reap(*os.waitpid(-1, 0))
 291         except OSError as e:
 292             if e.errno == errno.ECHILD: return
 293             raise
 294     def reap(self, pid, status):
 295         """Reaps a process."""
 296         # ooh, zombie process. reap it
 297         proc, args, python, on_success, on_error = self.running.pop(pid)
 298         # XXX: this is slightly dangerous; should actually use
 299         # temporary files
 300         stdout = proc.stdout.read()
 301         stderr = proc.stderr.read()
 302         self._log(stdout, stderr)
 303         if status:
 304             if python: eclass = PythonCallError
 305             else: eclass = CallError
 306             on_error(eclass(proc.returncode, args, stdout, stderr))
 307             return
 308         on_success(stdout, stderr)
 309
 310
 311 class DummyParallelShell(ParallelShell):
 312     """Same API as :class:`ParallelShell`, but doesn't actually
 313     parallelize (i.e. all calls to :meth:`wait` block.)"""
 314     def __init__(self, dry = False):
 315         super(DummyParallelShell, self).__init__(dry=dry, max=1)
 316
 317 class Error(wizard.Error):
 318     """Base exception for this module"""
 319     pass
 320
 321 class CallError(Error):
 322     """Indicates that a subprocess call returned a nonzero exit status."""
 323     #: The exit code of the failed subprocess.
 324     code = None
 325     #: List of the program and arguments that failed.
 326     args = None
 327     #: The stdout of the program.
 328     stdout = None
 329     #: The stderr of the program.
 330     stderr = None
 331     def __init__(self, code, args, stdout, stderr):
 332         self.code = code
 333         self.args = args
 334         self.stdout = stdout
 335         self.stderr = stderr
 336     def __str__(self):
 337         compact = self.stderr.rstrip().split("\n")[-1]
 338         return "%s (exited with %d)\n%s" % (compact, self.code, self.stderr)
 339
 340 class PythonCallError(CallError):
 341     """
 342     Indicates that a Python subprocess call had an uncaught exception.
 343     This exception also contains the attributes of :class:`CallError`.
 344     """
 345     #: Name of the uncaught exception.
 346     name = None
 347     def __init__(self, code, args, stdout, stderr):
 348         if stderr: self.name = util.get_exception_name(stderr)
 349         CallError.__init__(self, code, args, stdout, stderr)
 350     def __str__(self):
 351         if self.name:
 352             return "PythonCallError [%s]\n%s" % (self.name, self.stderr)
 353         else:
 354             return "PythonCallError\n%s" % self.stderr
 355