X-Git-Url: https://scripts.mit.edu/gitweb/wizard.git/blobdiff_plain/df934197496ca384a111c0cc10cd407e2368ffd2..319f82896fb658d9bbaf8d01d6c2c7410eb967c7:/wizard/shell.py diff --git a/wizard/shell.py b/wizard/shell.py index 79b2a71..b8d4bd1 100644 --- a/wizard/shell.py +++ b/wizard/shell.py @@ -25,12 +25,14 @@ 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, options): +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. + 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 @@ -43,7 +45,7 @@ def drop_priviledges(dir, options): args.append("%s=%s" % (k,v)) args += sys.argv logging.debug("Dropping priviledges") - if options.log_file: os.chown(options.log_file, uid, -1) + if log_file: os.chown(log_file, uid, -1) os.execlp('sudo', 'sudo', '-u', '#' + str(uid), *args) class Shell(object): @@ -53,6 +55,7 @@ class Shell(object): """ def __init__(self, dry = False): self.dry = dry + self.cwd = None def call(self, *args, **kwargs): """ Performs a system call. The actual executable and options should @@ -73,6 +76,9 @@ class Shell(object): 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. @@ -82,12 +88,14 @@ class Shell(object): >>> sh.call("cat", input='Foobar') ('Foobar', '') """ - if hasattr(self, "_wait"): - self._wait() + self._wait() kwargs.setdefault("interactive", False) kwargs.setdefault("strip", False) kwargs.setdefault("python", None) 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) @@ -108,9 +116,9 @@ class Shell(object): 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 @@ -120,11 +128,16 @@ class Shell(object): # ourself, and then setting up a # SIGCHILD handler to write a single byte to the pipe to get # us out of select() when a subprocess exits. - proc = subprocess.Popen(args, stdout=stdout, stderr=stderr, stdin=stdin) - if hasattr(self, "_async"): - self._async(proc, args, **kwargs) + proc = subprocess.Popen(args, stdout=stdout, stderr=stderr, stdin=stdin, cwd=self.cwd, ) + if self._async(proc, args, **kwargs): return proc stdout, stderr = proc.communicate(kwargs["input"]) + # can occur if we were doing interactive communication; i.e. + # we didn't pass in PIPE. + if stdout is None: + stdout = "" + if stderr is None: + stderr = "" if not kwargs["interactive"]: if kwargs["strip"]: self._log(None, stderr) @@ -135,7 +148,7 @@ class Shell(object): else: eclass = CallError raise eclass(proc.returncode, args, stdout, stderr) if kwargs["strip"]: - return stdout.rstrip("\n") + return str(stdout).rstrip("\n") return (stdout, stderr) def _log(self, stdout, stderr): """Logs the standard output and standard input from a command.""" @@ -143,6 +156,10 @@ class Shell(object): logging.debug("STDOUT:\n" + stdout) if stderr: logging.debug("STDERR:\n" + stderr) + def _wait(self): + pass + def _async(self, *args, **kwargs): + return False def callAsUser(self, *args, **kwargs): """ Performs a system call as a different user. This is only possible @@ -178,6 +195,8 @@ class Shell(object): on working directory context. Keyword arguments are the same as :meth:`call`. """ + if os.getuid(): + return self.call(*args, **kwargs) uid = os.stat(os.getcwd()).st_uid # consider also checking ruid? if uid != os.geteuid(): @@ -197,6 +216,12 @@ class Shell(object): """ kwargs["strip"] = True return self.call(*args, **kwargs) + def setcwd(self, cwd): + """ + Sets the directory processes are executed in. This sets a value + to be passed as the ``cwd`` argument to ``subprocess.Popen``. + """ + self.cwd = cwd class ParallelShell(Shell): """ @@ -245,7 +270,7 @@ class ParallelShell(Shell): self.running = {} self.max = max # maximum of commands to run in parallel @staticmethod - def make(self, no_parallelize, max): + def make(no_parallelize, max): """Convenience method oriented towards command modules.""" if no_parallelize: return DummyParallelShell() @@ -257,12 +282,13 @@ class ParallelShell(Shell): execution. See :meth:`Shell.call` source code for details. """ self.running[proc.pid] = (proc, args, python, on_success, on_error) + return True # so that the parent function returns def _wait(self): """ 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 + # XXX: This API sucks; the actual call/callAsUser call should # probably block automatically (unless I have a good reason not to) # bail out immediately on initial ramp up if len(self.running) < self.max: return @@ -296,6 +322,12 @@ class ParallelShell(Shell): return on_success(stdout, stderr) +# Setup a convenience global instance +shell = Shell() +call = shell.call +callAsUser = shell.callAsUser +safeCall = shell.safeCall +eval = shell.eval class DummyParallelShell(ParallelShell): """Same API as :class:`ParallelShell`, but doesn't actually