]> scripts.mit.edu Git - wizard.git/blobdiff - wizard/shell.py
Properly register salt values as random.
[wizard.git] / wizard / shell.py
index bdaa02ecfade635d2cea06d2309e8732dda0e902..b8d4bd1e6dcfa8cda550c192b63a74a3146a435c 100644 (file)
@@ -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"
 
     """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
     """
     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
     """
     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")
             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):
     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
     """
     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
     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``.
             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.
 
         :returns: a tuple of strings ``(stdout, stderr)``, or a string ``stdout``
             if ``strip`` is specified.
 
@@ -82,10 +88,14 @@ class Shell(object):
         >>> sh.call("cat", input='Foobar')
         ('Foobar', '')
         """
         >>> sh.call("cat", input='Foobar')
         ('Foobar', '')
         """
+        self._wait()
         kwargs.setdefault("interactive", False)
         kwargs.setdefault("strip", False)
         kwargs.setdefault("python", None)
         kwargs.setdefault("log", None)
         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)
         msg = "Running `" + ' '.join(args) + "`"
         if kwargs["strip"] and not kwargs["log"] is True or kwargs["log"] is False:
             logging.debug(msg)
@@ -106,9 +116,9 @@ class Shell(object):
             stdin=sys.stdin
             stderr=sys.stderr
         else:
             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
         # 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
@@ -118,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.
         # 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"])
             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)
         if not kwargs["interactive"]:
             if kwargs["strip"]:
                 self._log(None, stderr)
@@ -133,7 +148,7 @@ class Shell(object):
             else: eclass = CallError
             raise eclass(proc.returncode, args, stdout, stderr)
         if kwargs["strip"]:
             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."""
         return (stdout, stderr)
     def _log(self, stdout, stderr):
         """Logs the standard output and standard input from a command."""
@@ -141,6 +156,10 @@ class Shell(object):
             logging.debug("STDOUT:\n" + stdout)
         if stderr:
             logging.debug("STDERR:\n" + stderr)
             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
     def callAsUser(self, *args, **kwargs):
         """
         Performs a system call as a different user.  This is only possible
@@ -176,6 +195,8 @@ class Shell(object):
         on working directory context.  Keyword arguments are the
         same as :meth:`call`.
         """
         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():
         uid = os.stat(os.getcwd()).st_uid
         # consider also checking ruid?
         if uid != os.geteuid():
@@ -195,6 +216,12 @@ class Shell(object):
         """
         kwargs["strip"] = True
         return self.call(*args, **kwargs)
         """
         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):
     """
 
 class ParallelShell(Shell):
     """
@@ -203,19 +230,10 @@ class ParallelShell(Shell):
     with ``max`` subprocesses, and result in callback execution
     when they finish.
 
     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)
 
     .. 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).
         are the same as :meth:`Shell.call` with the following additions:
 
         :param on_success: Callback function for success (zero exit status).
@@ -251,22 +269,26 @@ class ParallelShell(Shell):
         super(ParallelShell, self).__init__(dry=dry)
         self.running = {}
         self.max = max # maximum of commands to run in parallel
         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 _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
         # 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
@@ -300,6 +322,12 @@ class ParallelShell(Shell):
             return
         on_success(stdout, stderr)
 
             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
 
 class DummyParallelShell(ParallelShell):
     """Same API as :class:`ParallelShell`, but doesn't actually