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