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(dir, log_file):
30 Checks if we are running as root. If we are, attempt to drop
31 priviledges to the user who owns ``dir``, by re-calling
32 itself using sudo with exec, such that the new process subsumes our
33 current one. If ``log_file`` is passed, the file is chown'ed
34 to the user we are dropping priviledges to, so the subprocess
39 uid = util.get_dir_uid(dir)
43 for k,v in os.environ.items():
44 if k.startswith('WIZARD_') or k == "SSH_GSSAPI_NAME":
45 args.append("%s=%s" % (k,v))
47 logging.debug("Dropping priviledges")
48 if log_file: os.chown(log_file, uid, -1)
49 os.execlp('sudo', 'sudo', '-u', '#' + str(uid), *args)
53 An advanced shell that performs logging. If ``dry`` is ``True``,
54 no commands are actually run.
56 def __init__(self, dry = False):
59 def call(self, *args, **kwargs):
61 Performs a system call. The actual executable and options should
62 be passed as arguments to this function. It will magically
63 ensure that 'wizard' as a command works. Several keyword arguments
66 :param python: explicitly marks the subprocess as Python or not Python
67 for improved error reporting. By default, we use
68 :func:`is_python` to autodetect this.
69 :param input: input to feed the subprocess on standard input.
70 :param interactive: whether or not directly hook up all pipes
71 to the controlling terminal, to allow interaction with subprocess.
72 :param strip: if ``True``, instead of returning a tuple,
73 return the string stdout output of the command with trailing newlines
74 removed. This emulates the behavior of backticks and ``$()`` in Bash.
75 Prefer to use :meth:`eval` instead (you should only need to explicitly
76 specify this if you are using another wrapper around this function).
77 :param log: if True, we log the call as INFO, if False, we log the call
78 as DEBUG, otherwise, we detect based on ``strip``.
81 :param stdin: a file-type object that will be written to or read from as a pipe.
82 :returns: a tuple of strings ``(stdout, stderr)``, or a string ``stdout``
83 if ``strip`` is specified.
86 >>> sh.call("echo", "Foobar")
88 >>> sh.call("cat", input='Foobar')
92 kwargs.setdefault("interactive", False)
93 kwargs.setdefault("strip", False)
94 kwargs.setdefault("python", None)
95 kwargs.setdefault("log", None)
96 kwargs.setdefault("stdout", subprocess.PIPE)
97 kwargs.setdefault("stdin", subprocess.PIPE)
98 kwargs.setdefault("stderr", subprocess.PIPE)
99 msg = "Running `" + ' '.join(args) + "`"
100 if kwargs["strip"] and not kwargs["log"] is True or kwargs["log"] is False:
108 if kwargs["python"] is None and is_python(args):
109 kwargs["python"] = True
110 if args[0] == "wizard":
113 kwargs.setdefault("input", None)
114 if kwargs["interactive"]:
119 stdout=kwargs["stdout"]
120 stdin=kwargs["stdin"]
121 stderr=kwargs["stderr"]
122 # XXX: There is a possible problem here where we can fill up
123 # the kernel buffer if we have 64KB of data. This shouldn't
124 # be a problem, and the fix for such case would be to write to
125 # temporary files instead of a pipe.
126 # Another possible way of fixing this is converting from a
127 # waitpid() pump to a select() pump, creating a pipe to
128 # ourself, and then setting up a
129 # SIGCHILD handler to write a single byte to the pipe to get
130 # us out of select() when a subprocess exits.
131 proc = subprocess.Popen(args, stdout=stdout, stderr=stderr, stdin=stdin, cwd=self.cwd, )
132 if self._async(proc, args, **kwargs):
134 stdout, stderr = proc.communicate(kwargs["input"])
135 # can occur if we were doing interactive communication; i.e.
136 # we didn't pass in PIPE.
141 if not kwargs["interactive"]:
143 self._log(None, stderr)
145 self._log(stdout, stderr)
147 if kwargs["python"]: eclass = PythonCallError
148 else: eclass = CallError
149 raise eclass(proc.returncode, args, stdout, stderr)
151 return str(stdout).rstrip("\n")
152 return (stdout, stderr)
153 def _log(self, stdout, stderr):
154 """Logs the standard output and standard input from a command."""
156 logging.debug("STDOUT:\n" + stdout)
158 logging.debug("STDERR:\n" + stderr)
161 def _async(self, *args, **kwargs):
163 def callAsUser(self, *args, **kwargs):
165 Performs a system call as a different user. This is only possible
166 if you are running as root. Keyword arguments
167 are the same as :meth:`call` with the following additions:
169 :param user: name of the user to run command as.
170 :param uid: uid of the user to run command as.
174 The resulting system call internally uses :command:`sudo`,
175 and as such environment variables will get scrubbed. We
176 manually preserve :envvar:`SSH_GSSAPI_NAME`.
178 user = kwargs.pop("user", None)
179 uid = kwargs.pop("uid", None)
180 kwargs.setdefault("python", is_python(args))
181 if not user and not uid: return self.call(*args, **kwargs)
182 if os.getenv("SSH_GSSAPI_NAME"):
183 # This might be generalized as "preserve some environment"
185 args.insert(0, "SSH_GSSAPI_NAME=" + os.getenv("SSH_GSSAPI_NAME"))
186 if uid: return self.call("sudo", "-u", "#" + str(uid), *args, **kwargs)
187 if user: return self.call("sudo", "-u", user, *args, **kwargs)
188 def safeCall(self, *args, **kwargs):
190 Checks if the owner of the current working directory is the same
191 as the current user, and if it isn't, attempts to sudo to be
192 that user. The intended use case is for calling Git commands
193 when running as root, but this method should be used when
194 interfacing with any moderately complex program that depends
195 on working directory context. Keyword arguments are the
196 same as :meth:`call`.
199 return self.call(*args, **kwargs)
200 uid = os.stat(os.getcwd()).st_uid
201 # consider also checking ruid?
202 if uid != os.geteuid():
204 return self.callAsUser(*args, **kwargs)
206 return self.call(*args, **kwargs)
207 def eval(self, *args, **kwargs):
209 Evaluates a command and returns its output, with trailing newlines
210 stripped (like backticks in Bash). This is a convenience method for
211 calling :meth:`call` with ``strip``.
214 >>> sh.eval("echo", "Foobar")
217 kwargs["strip"] = True
218 return self.call(*args, **kwargs)
219 def setcwd(self, cwd):
221 Sets the directory processes are executed in. This sets a value
222 to be passed as the ``cwd`` argument to ``subprocess.Popen``.
226 class ParallelShell(Shell):
228 Modifies the semantics of :class:`Shell` so that
229 commands are queued here, and executed in parallel using waitpid
230 with ``max`` subprocesses, and result in callback execution
233 .. method:: call(*args, **kwargs)
235 Enqueues a system call for parallel processing. If there are
236 no openings in the queue, this will block. Keyword arguments
237 are the same as :meth:`Shell.call` with the following additions:
239 :param on_success: Callback function for success (zero exit status).
240 The callback function should accept two arguments,
241 ``stdout`` and ``stderr``.
242 :param on_error: Callback function for failure (nonzero exit status).
243 The callback function should accept one argument, the
244 exception that would have been thrown by the synchronous
246 :return: The :class:`subprocess.Proc` object that was opened.
248 .. method:: callAsUser(*args, **kwargs)
250 Enqueues a system call under a different user for parallel
251 processing. Keyword arguments are the same as
252 :meth:`Shell.callAsUser` with the additions of keyword
253 arguments from :meth:`call`.
255 .. method:: safeCall(*args, **kwargs)
257 Enqueues a "safe" call for parallel processing. Keyword
258 arguments are the same as :meth:`Shell.safeCall` with the
259 additions of keyword arguments from :meth:`call`.
261 .. method:: eval(*args, **kwargs)
263 No difference from :meth:`call`. Consider having a
264 non-parallel shell if the program you are shelling out
268 def __init__(self, dry = False, max = 10):
269 super(ParallelShell, self).__init__(dry=dry)
271 self.max = max # maximum of commands to run in parallel
273 def make(no_parallelize, max):
274 """Convenience method oriented towards command modules."""
276 return DummyParallelShell()
278 return ParallelShell(max=max)
279 def _async(self, proc, args, python, on_success, on_error, **kwargs):
281 Gets handed a :class:`subprocess.Proc` object from our deferred
282 execution. See :meth:`Shell.call` source code for details.
284 self.running[proc.pid] = (proc, args, python, on_success, on_error)
285 return True # so that the parent function returns
288 Blocking call that waits for an open subprocess slot. This is
289 automatically called by :meth:`Shell.call`.
291 # XXX: This API sucks; the actual call/callAsUser call should
292 # probably block automatically (unless I have a good reason not to)
293 # bail out immediately on initial ramp up
294 if len(self.running) < self.max: return
295 # now, wait for open pids.
297 self.reap(*os.waitpid(-1, 0))
299 if e.errno == errno.ECHILD: return
302 """Waits for all of our subprocesses to terminate."""
305 self.reap(*os.waitpid(-1, 0))
307 if e.errno == errno.ECHILD: return
309 def reap(self, pid, status):
310 """Reaps a process."""
311 # ooh, zombie process. reap it
312 proc, args, python, on_success, on_error = self.running.pop(pid)
313 # XXX: this is slightly dangerous; should actually use
315 stdout = proc.stdout.read()
316 stderr = proc.stderr.read()
317 self._log(stdout, stderr)
319 if python: eclass = PythonCallError
320 else: eclass = CallError
321 on_error(eclass(proc.returncode, args, stdout, stderr))
323 on_success(stdout, stderr)
325 # Setup a convenience global instance
328 callAsUser = shell.callAsUser
329 safeCall = shell.safeCall
332 class DummyParallelShell(ParallelShell):
333 """Same API as :class:`ParallelShell`, but doesn't actually
334 parallelize (i.e. all calls to :meth:`wait` block.)"""
335 def __init__(self, dry = False):
336 super(DummyParallelShell, self).__init__(dry=dry, max=1)
338 class Error(wizard.Error):
339 """Base exception for this module"""
342 class CallError(Error):
343 """Indicates that a subprocess call returned a nonzero exit status."""
344 #: The exit code of the failed subprocess.
346 #: List of the program and arguments that failed.
348 #: The stdout of the program.
350 #: The stderr of the program.
352 def __init__(self, code, args, stdout, stderr):
358 compact = self.stderr.rstrip().split("\n")[-1]
359 return "%s (exited with %d)\n%s" % (compact, self.code, self.stderr)
361 class PythonCallError(CallError):
363 Indicates that a Python subprocess call had an uncaught exception.
364 This exception also contains the attributes of :class:`CallError`.
366 #: Name of the uncaught exception.
368 def __init__(self, code, args, stdout, stderr):
369 if stderr: self.name = util.get_exception_name(stderr)
370 CallError.__init__(self, code, args, stdout, stderr)
373 return "PythonCallError [%s]\n%s" % (self.name, self.stderr)
375 return "PythonCallError\n%s" % self.stderr