X-Git-Url: https://scripts.mit.edu/gitweb/wizard.git/blobdiff_plain/de4349af65d013973ddcf44a7565326d411e50eb..8264ae4501b4e3ca5e7fef548bebb33c16a62a65:/wizard/shell.py diff --git a/wizard/shell.py b/wizard/shell.py index 39cd102..8660d1a 100644 --- a/wizard/shell.py +++ b/wizard/shell.py @@ -10,6 +10,7 @@ import subprocess import logging import sys import os +import errno import wizard from wizard import util @@ -24,6 +25,29 @@ def is_python(args): """Detects whether or not an argument list invokes a Python program.""" return args[0] == "python" or args[0] == "wizard" +def drop_priviledges(dir, log_file): + """ + Checks if we are running as root. If we are, attempt to drop + priviledges to the user who owns ``dir``, by re-calling + itself using sudo with exec, such that the new process subsumes our + current one. If ``log_file`` is passed, the file is chown'ed + to the user we are dropping priviledges to, so the subprocess + can write to it. + """ + if os.getuid(): + return + uid = util.get_dir_uid(dir) + if not uid: + return + args = [] + for k,v in os.environ.items(): + if k.startswith('WIZARD_') or k == "SSH_GSSAPI_NAME": + args.append("%s=%s" % (k,v)) + args += sys.argv + logging.debug("Dropping priviledges") + if log_file: os.chown(log_file, uid, -1) + os.execlp('sudo', 'sudo', '-u', '#' + str(uid), *args) + class Shell(object): """ An advanced shell that performs logging. If ``dry`` is ``True``, @@ -49,6 +73,11 @@ class Shell(object): removed. This emulates the behavior of backticks and ``$()`` in Bash. Prefer to use :meth:`eval` instead (you should only need to explicitly specify this if you are using another wrapper around this function). + :param log: if True, we log the call as INFO, if False, we log the call + as DEBUG, otherwise, we detect based on ``strip``. + :param stdout: + :param stderr: + :param stdin: a file-type object that will be written to or read from as a pipe. :returns: a tuple of strings ``(stdout, stderr)``, or a string ``stdout`` if ``strip`` is specified. @@ -58,26 +87,38 @@ class Shell(object): >>> sh.call("cat", input='Foobar') ('Foobar', '') """ + if hasattr(self, "_wait"): + self._wait() + kwargs.setdefault("interactive", False) + kwargs.setdefault("strip", False) kwargs.setdefault("python", None) - logging.info("Running `" + ' '.join(args) + "`") + kwargs.setdefault("log", None) + kwargs.setdefault("stdout", subprocess.PIPE) + kwargs.setdefault("stdin", subprocess.PIPE) + kwargs.setdefault("stderr", subprocess.PIPE) + msg = "Running `" + ' '.join(args) + "`" + if kwargs["strip"] and not kwargs["log"] is True or kwargs["log"] is False: + logging.debug(msg) + else: + logging.info(msg) if self.dry: - return + if kwargs["strip"]: + return '' + return None, None if kwargs["python"] is None and is_python(args): kwargs["python"] = True if args[0] == "wizard": args = list(args) args[0] = wizard_bin kwargs.setdefault("input", None) - kwargs.setdefault("interactive", False) - kwargs.setdefault("strip", False) if kwargs["interactive"]: stdout=sys.stdout stdin=sys.stdin stderr=sys.stderr else: - stdout=subprocess.PIPE - stdin=subprocess.PIPE - stderr=subprocess.PIPE + stdout=kwargs["stdout"] + stdin=kwargs["stdin"] + stderr=kwargs["stderr"] # XXX: There is a possible problem here where we can fill up # the kernel buffer if we have 64KB of data. This shouldn't # be a problem, and the fix for such case would be to write to @@ -93,7 +134,10 @@ class Shell(object): return proc stdout, stderr = proc.communicate(kwargs["input"]) if not kwargs["interactive"]: - self._log(stdout, stderr) + if kwargs["strip"]: + self._log(None, stderr) + else: + self._log(stdout, stderr) if proc.returncode: if kwargs["python"]: eclass = PythonCallError else: eclass = CallError @@ -128,6 +172,7 @@ class Shell(object): if not user and not uid: return self.call(*args, **kwargs) if os.getenv("SSH_GSSAPI_NAME"): # This might be generalized as "preserve some environment" + args = list(args) args.insert(0, "SSH_GSSAPI_NAME=" + os.getenv("SSH_GSSAPI_NAME")) if uid: return self.call("sudo", "-u", "#" + str(uid), *args, **kwargs) if user: return self.call("sudo", "-u", user, *args, **kwargs) @@ -168,19 +213,10 @@ class ParallelShell(Shell): with ``max`` subprocesses, and result in callback execution when they finish. - Before enqueueing a system call with :meth:`call` or :meth:`callAsUser`, - you should wait for an open slot using :meth:`wait`; otherwise, - ``max`` rate limiting will have no effect. For example:: - - sh = ParallelShell() - for command in commands_to_execute_in_parallel: - sh.wait() - sh.call(*command) - sh.join() - .. method:: call(*args, **kwargs) - Enqueues a system call for parallel processing. Keyword arguments + Enqueues a system call for parallel processing. If there are + no openings in the queue, this will block. Keyword arguments are the same as :meth:`Shell.call` with the following additions: :param on_success: Callback function for success (zero exit status). @@ -216,20 +252,23 @@ class ParallelShell(Shell): super(ParallelShell, self).__init__(dry=dry) self.running = {} self.max = max # maximum of commands to run in parallel + @staticmethod + def make(no_parallelize, max): + """Convenience method oriented towards command modules.""" + if no_parallelize: + return DummyParallelShell() + else: + return ParallelShell(max=max) def _async(self, proc, args, python, on_success, on_error, **kwargs): """ Gets handed a :class:`subprocess.Proc` object from our deferred execution. See :meth:`Shell.call` source code for details. """ self.running[proc.pid] = (proc, args, python, on_success, on_error) - def wait(self): + def _wait(self): """ - Blocking call that waits for an open subprocess slot. You should - call this before enqueuing. - - .. note:: - - This method may become unnecessary in the future. + Blocking call that waits for an open subprocess slot. This is + automatically called by :meth:`Shell.call`. """ # XXX: This API sucks; the actuall call/callAsUser call should # probably block automatically (unless I have a good reason not to) @@ -237,10 +276,20 @@ class ParallelShell(Shell): if len(self.running) < self.max: return # now, wait for open pids. try: - pid, status = os.waitpid(-1, 0) + self.reap(*os.waitpid(-1, 0)) + except OSError as e: + if e.errno == errno.ECHILD: return + raise + def join(self): + """Waits for all of our subprocesses to terminate.""" + try: + while True: + self.reap(*os.waitpid(-1, 0)) except OSError as e: if e.errno == errno.ECHILD: return - raise e + raise + def reap(self, pid, status): + """Reaps a process.""" # ooh, zombie process. reap it proc, args, python, on_success, on_error = self.running.pop(pid) # XXX: this is slightly dangerous; should actually use @@ -254,14 +303,7 @@ class ParallelShell(Shell): on_error(eclass(proc.returncode, args, stdout, stderr)) return on_success(stdout, stderr) - def join(self): - """Waits for all of our subprocesses to terminate.""" - try: - while os.waitpid(-1, 0): - pass - except OSError as e: - if e.errno == errno.ECHILD: return - raise e + class DummyParallelShell(ParallelShell): """Same API as :class:`ParallelShell`, but doesn't actually @@ -289,7 +331,8 @@ class CallError(Error): self.stdout = stdout self.stderr = stderr def __str__(self): - return "CallError [%d]" % self.code + compact = self.stderr.rstrip().split("\n")[-1] + return "%s (exited with %d)\n%s" % (compact, self.code, self.stderr) class PythonCallError(CallError): """ @@ -303,7 +346,7 @@ class PythonCallError(CallError): CallError.__init__(self, code, args, stdout, stderr) def __str__(self): if self.name: - return "PythonCallError [%s]" % self.name + return "PythonCallError [%s]\n%s" % (self.name, self.stderr) else: - return "PythonCallError" + return "PythonCallError\n%s" % self.stderr