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