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