2 Wrappers around subprocess functionality that simulate an actual shell.
6 from wizard.shell import *
16 from wizard import util
18 wizard_bin = sys.argv[0]
20 This is the path to the wizard executable as specified
21 by the caller; it lets us recursively invoke wizard.
25 """Detects whether or not an argument list invokes a Python program."""
26 return args[0] == "python" or args[0] == "wizard"
28 def drop_priviledges(options):
30 Checks if we are running as root. If we are, attempt to drop
31 priviledges to the user who owns the current directory, by re-calling
32 itself using sudo with exec, such that the new process subsumes our
37 uid = util.get_dir_uid('.')
41 for k,v in os.environ.items():
42 if k.startswith('WIZARD_') or k == "SSH_GSSAPI_NAME":
43 args.append("%s=%s" % (k,v))
45 logging.debug("Dropping priviledges")
46 if options.log_file: os.chown(options.log_file, uid, -1)
47 os.execlp('sudo', 'sudo', '-u', '#' + str(uid), *args)
51 An advanced shell that performs logging. If ``dry`` is ``True``,
52 no commands are actually run.
54 def __init__(self, dry = False):
56 def call(self, *args, **kwargs):
58 Performs a system call. The actual executable and options should
59 be passed as arguments to this function. It will magically
60 ensure that 'wizard' as a command works. Several keyword arguments
63 :param python: explicitly marks the subprocess as Python or not Python
64 for improved error reporting. By default, we use
65 :func:`is_python` to autodetect this.
66 :param input: input to feed the subprocess on standard input.
67 :param interactive: whether or not directly hook up all pipes
68 to the controlling terminal, to allow interaction with subprocess.
69 :param strip: if ``True``, instead of returning a tuple,
70 return the string stdout output of the command with trailing newlines
71 removed. This emulates the behavior of backticks and ``$()`` in Bash.
72 Prefer to use :meth:`eval` instead (you should only need to explicitly
73 specify this if you are using another wrapper around this function).
74 :param log: if True, we log the call as INFO, if False, we log the call
75 as DEBUG, otherwise, we detect based on ``strip``.
76 :returns: a tuple of strings ``(stdout, stderr)``, or a string ``stdout``
77 if ``strip`` is specified.
80 >>> sh.call("echo", "Foobar")
82 >>> sh.call("cat", input='Foobar')
85 kwargs.setdefault("interactive", False)
86 kwargs.setdefault("strip", False)
87 kwargs.setdefault("python", None)
88 kwargs.setdefault("log", None)
89 msg = "Running `" + ' '.join(args) + "`"
90 if kwargs["strip"] and not kwargs["log"] is True or kwargs["log"] is False:
96 if kwargs["python"] is None and is_python(args):
97 kwargs["python"] = True
98 if args[0] == "wizard":
101 kwargs.setdefault("input", None)
102 if kwargs["interactive"]:
107 stdout=subprocess.PIPE
108 stdin=subprocess.PIPE
109 stderr=subprocess.PIPE
110 # XXX: There is a possible problem here where we can fill up
111 # the kernel buffer if we have 64KB of data. This shouldn't
112 # be a problem, and the fix for such case would be to write to
113 # temporary files instead of a pipe.
114 # Another possible way of fixing this is converting from a
115 # waitpid() pump to a select() pump, creating a pipe to
116 # ourself, and then setting up a
117 # SIGCHILD handler to write a single byte to the pipe to get
118 # us out of select() when a subprocess exits.
119 proc = subprocess.Popen(args, stdout=stdout, stderr=stderr, stdin=stdin)
120 if hasattr(self, "_async"):
121 self._async(proc, args, **kwargs)
123 stdout, stderr = proc.communicate(kwargs["input"])
124 if not kwargs["interactive"]:
126 self._log(None, stderr)
128 self._log(stdout, stderr)
130 if kwargs["python"]: eclass = PythonCallError
131 else: eclass = CallError
132 raise eclass(proc.returncode, args, stdout, stderr)
134 return stdout.rstrip("\n")
135 return (stdout, stderr)
136 def _log(self, stdout, stderr):
137 """Logs the standard output and standard input from a command."""
139 logging.debug("STDOUT:\n" + stdout)
141 logging.debug("STDERR:\n" + stderr)
142 def callAsUser(self, *args, **kwargs):
144 Performs a system call as a different user. This is only possible
145 if you are running as root. Keyword arguments
146 are the same as :meth:`call` with the following additions:
148 :param user: name of the user to run command as.
149 :param uid: uid of the user to run command as.
153 The resulting system call internally uses :command:`sudo`,
154 and as such environment variables will get scrubbed. We
155 manually preserve :envvar:`SSH_GSSAPI_NAME`.
157 user = kwargs.pop("user", None)
158 uid = kwargs.pop("uid", None)
159 kwargs.setdefault("python", is_python(args))
160 if not user and not uid: return self.call(*args, **kwargs)
161 if os.getenv("SSH_GSSAPI_NAME"):
162 # This might be generalized as "preserve some environment"
164 args.insert(0, "SSH_GSSAPI_NAME=" + os.getenv("SSH_GSSAPI_NAME"))
165 if uid: return self.call("sudo", "-u", "#" + str(uid), *args, **kwargs)
166 if user: return self.call("sudo", "-u", user, *args, **kwargs)
167 def safeCall(self, *args, **kwargs):
169 Checks if the owner of the current working directory is the same
170 as the current user, and if it isn't, attempts to sudo to be
171 that user. The intended use case is for calling Git commands
172 when running as root, but this method should be used when
173 interfacing with any moderately complex program that depends
174 on working directory context. Keyword arguments are the
175 same as :meth:`call`.
177 uid = os.stat(os.getcwd()).st_uid
178 # consider also checking ruid?
179 if uid != os.geteuid():
181 return self.callAsUser(*args, **kwargs)
183 return self.call(*args, **kwargs)
184 def eval(self, *args, **kwargs):
186 Evaluates a command and returns its output, with trailing newlines
187 stripped (like backticks in Bash). This is a convenience method for
188 calling :meth:`call` with ``strip``.
191 >>> sh.eval("echo", "Foobar")
194 kwargs["strip"] = True
195 return self.call(*args, **kwargs)
197 class ParallelShell(Shell):
199 Modifies the semantics of :class:`Shell` so that
200 commands are queued here, and executed in parallel using waitpid
201 with ``max`` subprocesses, and result in callback execution
204 Before enqueueing a system call with :meth:`call` or :meth:`callAsUser`,
205 you should wait for an open slot using :meth:`wait`; otherwise,
206 ``max`` rate limiting will have no effect. For example::
209 for command in commands_to_execute_in_parallel:
214 .. method:: call(*args, **kwargs)
216 Enqueues a system call for parallel processing. Keyword arguments
217 are the same as :meth:`Shell.call` with the following additions:
219 :param on_success: Callback function for success (zero exit status).
220 The callback function should accept two arguments,
221 ``stdout`` and ``stderr``.
222 :param on_error: Callback function for failure (nonzero exit status).
223 The callback function should accept one argument, the
224 exception that would have been thrown by the synchronous
226 :return: The :class:`subprocess.Proc` object that was opened.
228 .. method:: callAsUser(*args, **kwargs)
230 Enqueues a system call under a different user for parallel
231 processing. Keyword arguments are the same as
232 :meth:`Shell.callAsUser` with the additions of keyword
233 arguments from :meth:`call`.
235 .. method:: safeCall(*args, **kwargs)
237 Enqueues a "safe" call for parallel processing. Keyword
238 arguments are the same as :meth:`Shell.safeCall` with the
239 additions of keyword arguments from :meth:`call`.
241 .. method:: eval(*args, **kwargs)
243 No difference from :meth:`call`. Consider having a
244 non-parallel shell if the program you are shelling out
248 def __init__(self, dry = False, max = 10):
249 super(ParallelShell, self).__init__(dry=dry)
251 self.max = max # maximum of commands to run in parallel
252 def _async(self, proc, args, python, on_success, on_error, **kwargs):
254 Gets handed a :class:`subprocess.Proc` object from our deferred
255 execution. See :meth:`Shell.call` source code for details.
257 self.running[proc.pid] = (proc, args, python, on_success, on_error)
260 Blocking call that waits for an open subprocess slot. You should
261 call this before enqueuing.
265 This method may become unnecessary in the future.
267 # XXX: This API sucks; the actuall call/callAsUser call should
268 # probably block automatically (unless I have a good reason not to)
269 # bail out immediately on initial ramp up
270 if len(self.running) < self.max: return
271 # now, wait for open pids.
273 self.reap(*os.waitpid(-1, 0))
275 if e.errno == errno.ECHILD: return
278 """Waits for all of our subprocesses to terminate."""
281 self.reap(*os.waitpid(-1, 0))
283 if e.errno == errno.ECHILD: return
285 def reap(self, pid, status):
286 """Reaps a process."""
287 # ooh, zombie process. reap it
288 proc, args, python, on_success, on_error = self.running.pop(pid)
289 # XXX: this is slightly dangerous; should actually use
291 stdout = proc.stdout.read()
292 stderr = proc.stderr.read()
293 self._log(stdout, stderr)
295 if python: eclass = PythonCallError
296 else: eclass = CallError
297 on_error(eclass(proc.returncode, args, stdout, stderr))
299 on_success(stdout, stderr)
302 class DummyParallelShell(ParallelShell):
303 """Same API as :class:`ParallelShell`, but doesn't actually
304 parallelize (i.e. all calls to :meth:`wait` block.)"""
305 def __init__(self, dry = False):
306 super(DummyParallelShell, self).__init__(dry=dry, max=1)
308 class Error(wizard.Error):
309 """Base exception for this module"""
312 class CallError(Error):
313 """Indicates that a subprocess call returned a nonzero exit status."""
314 #: The exit code of the failed subprocess.
316 #: List of the program and arguments that failed.
318 #: The stdout of the program.
320 #: The stderr of the program.
322 def __init__(self, code, args, stdout, stderr):
328 return "CallError [%d]\n%s" % (self.code, self.stderr)
330 class PythonCallError(CallError):
332 Indicates that a Python subprocess call had an uncaught exception.
333 This exception also contains the attributes of :class:`CallError`.
335 #: Name of the uncaught exception.
337 def __init__(self, code, args, stdout, stderr):
338 if stderr: self.name = util.get_exception_name(stderr)
339 CallError.__init__(self, code, args, stdout, stderr)
342 return "PythonCallError [%s]\n%s" % (self.name, self.stderr)
344 return "PythonCallError\n%s" % self.stderr