"""Detects whether or not an argument list invokes a Python program."""
return args[0] == "python" or args[0] == "wizard"
-def drop_priviledges(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 the current directory, by re-calling
+ 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
- uid = util.get_dir_uid('.')
+ uid = util.get_dir_uid(dir)
if not uid:
return
args = []
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):
"""
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
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 addenv: mapping of environment variables *to add*
+ :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.
>>> sh.call("cat", input='Foobar')
('Foobar', '')
"""
+ self._wait()
kwargs.setdefault("interactive", False)
kwargs.setdefault("strip", False)
kwargs.setdefault("python", None)
- if kwargs["strip"]:
- logging.debug("Evaluating `" + ' '.join(args) + "`")
+ kwargs.setdefault("log", None)
+ kwargs.setdefault("stdout", subprocess.PIPE)
+ kwargs.setdefault("stdin", subprocess.PIPE)
+ kwargs.setdefault("stderr", subprocess.PIPE)
+ kwargs.setdefault("addenv", None)
+ msg = "Running `" + ' '.join(args) + "`"
+ if kwargs["strip"] and not kwargs["log"] is True or kwargs["log"] is False:
+ logging.debug(msg)
else:
- logging.info("Running `" + ' '.join(args) + "`")
+ 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":
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"]
+ env = None
+ if kwargs["addenv"]:
+ env = dict(os.environ.items() + kwargs["addenv"].items())
# 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
+ # normally be a problem, and the fix for such case would be to write to
# temporary files instead of a pipe.
+ #
+ # However, it *is* a problem when you do something silly, like
+ # pass --debug to mass-upgrade.
+ #
# Another possible way of fixing this is converting from a
# waitpid() pump to a select() pump, creating a pipe to
- # 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)
+ # 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, cwd=self.cwd, env=env)
+ 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)
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."""
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
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():
"""
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
+ def interactive():
+ user_shell = os.getenv("SHELL")
+ if not user_shell: user_shell = "/bin/bash"
+ # XXX: scripts specific hack, since mbash doesn't respect the current working directory
+ # When the revolution comes (i.e. $ATHENA_HOMEDIR/Scripts is your Scripts home
+ # directory) this isn't strictly necessary, but we'll probably need to support
+ # web_scripts directories ad infinitum.
+ if user_shell == "/usr/local/bin/mbash": user_shell = "/bin/bash"
+
+ try:
+ self.call(user_shell, "-i", interactive=True)
+ except shell.CallError as e:
+ logging.warning("Shell returned non-zero exit code %d" % e.code)
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).
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):
+ return True # so that the parent function returns
+ 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
+ # 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
on_error(eclass(proc.returncode, args, stdout, stderr))
return
on_success(stdout, stderr)
+ def interactive():
+ raise Error("Cannot use interactive() on parallel shell")
+# Setup a convenience global instance
+shell = Shell()
+call = shell.call
+callAsUser = shell.callAsUser
+safeCall = shell.safeCall
+eval = shell.eval
+interactive = shell.interactive
class DummyParallelShell(ParallelShell):
"""Same API as :class:`ParallelShell`, but doesn't actually
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):
"""
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
scripts.mit.edu / wizard.git / blobdiff