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):
58 def call(self, *args, **kwargs):
60 Performs a system call. The actual executable and options should
61 be passed as arguments to this function. It will magically
62 ensure that 'wizard' as a command works. Several keyword arguments
65 :param python: explicitly marks the subprocess as Python or not Python
66 for improved error reporting. By default, we use
67 :func:`is_python` to autodetect this.
68 :param input: input to feed the subprocess on standard input.
69 :param interactive: whether or not directly hook up all pipes
70 to the controlling terminal, to allow interaction with subprocess.
71 :param strip: if ``True``, instead of returning a tuple,
72 return the string stdout output of the command with trailing newlines
73 removed. This emulates the behavior of backticks and ``$()`` in Bash.
74 Prefer to use :meth:`eval` instead (you should only need to explicitly
75 specify this if you are using another wrapper around this function).
76 :param log: if True, we log the call as INFO, if False, we log the call
77 as DEBUG, otherwise, we detect based on ``strip``.
80 :param stdin: a file-type object that will be written to or read from as a pipe.
81 :returns: a tuple of strings ``(stdout, stderr)``, or a string ``stdout``
82 if ``strip`` is specified.
85 >>> sh.call("echo", "Foobar")
87 >>> sh.call("cat", input='Foobar')
90 if hasattr(self, "_wait"):
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)
132 if hasattr(self, "_async"):
133 self._async(proc, args, **kwargs)
135 stdout, stderr = proc.communicate(kwargs["input"])
136 if not kwargs["interactive"]:
138 self._log(None, stderr)
140 self._log(stdout, stderr)
142 if kwargs["python"]: eclass = PythonCallError
143 else: eclass = CallError
144 raise eclass(proc.returncode, args, stdout, stderr)
146 return stdout.rstrip("\n")
147 return (stdout, stderr)
148 def _log(self, stdout, stderr):
149 """Logs the standard output and standard input from a command."""
151 logging.debug("STDOUT:\n" + stdout)
153 logging.debug("STDERR:\n" + stderr)
154 def callAsUser(self, *args, **kwargs):
156 Performs a system call as a different user. This is only possible
157 if you are running as root. Keyword arguments
158 are the same as :meth:`call` with the following additions:
160 :param user: name of the user to run command as.
161 :param uid: uid of the user to run command as.
165 The resulting system call internally uses :command:`sudo`,
166 and as such environment variables will get scrubbed. We
167 manually preserve :envvar:`SSH_GSSAPI_NAME`.
169 user = kwargs.pop("user", None)
170 uid = kwargs.pop("uid", None)
171 kwargs.setdefault("python", is_python(args))
172 if not user and not uid: return self.call(*args, **kwargs)
173 if os.getenv("SSH_GSSAPI_NAME"):
174 # This might be generalized as "preserve some environment"
176 args.insert(0, "SSH_GSSAPI_NAME=" + os.getenv("SSH_GSSAPI_NAME"))
177 if uid: return self.call("sudo", "-u", "#" + str(uid), *args, **kwargs)
178 if user: return self.call("sudo", "-u", user, *args, **kwargs)
179 def safeCall(self, *args, **kwargs):
181 Checks if the owner of the current working directory is the same
182 as the current user, and if it isn't, attempts to sudo to be
183 that user. The intended use case is for calling Git commands
184 when running as root, but this method should be used when
185 interfacing with any moderately complex program that depends
186 on working directory context. Keyword arguments are the
187 same as :meth:`call`.
189 uid = os.stat(os.getcwd()).st_uid
190 # consider also checking ruid?
191 if uid != os.geteuid():
193 return self.callAsUser(*args, **kwargs)
195 return self.call(*args, **kwargs)
196 def eval(self, *args, **kwargs):
198 Evaluates a command and returns its output, with trailing newlines
199 stripped (like backticks in Bash). This is a convenience method for
200 calling :meth:`call` with ``strip``.
203 >>> sh.eval("echo", "Foobar")
206 kwargs["strip"] = True
207 return self.call(*args, **kwargs)
209 class ParallelShell(Shell):
211 Modifies the semantics of :class:`Shell` so that
212 commands are queued here, and executed in parallel using waitpid
213 with ``max`` subprocesses, and result in callback execution
216 .. method:: call(*args, **kwargs)
218 Enqueues a system call for parallel processing. If there are
219 no openings in the queue, this will block. Keyword arguments
220 are the same as :meth:`Shell.call` with the following additions:
222 :param on_success: Callback function for success (zero exit status).
223 The callback function should accept two arguments,
224 ``stdout`` and ``stderr``.
225 :param on_error: Callback function for failure (nonzero exit status).
226 The callback function should accept one argument, the
227 exception that would have been thrown by the synchronous
229 :return: The :class:`subprocess.Proc` object that was opened.
231 .. method:: callAsUser(*args, **kwargs)
233 Enqueues a system call under a different user for parallel
234 processing. Keyword arguments are the same as
235 :meth:`Shell.callAsUser` with the additions of keyword
236 arguments from :meth:`call`.
238 .. method:: safeCall(*args, **kwargs)
240 Enqueues a "safe" call for parallel processing. Keyword
241 arguments are the same as :meth:`Shell.safeCall` with the
242 additions of keyword arguments from :meth:`call`.
244 .. method:: eval(*args, **kwargs)
246 No difference from :meth:`call`. Consider having a
247 non-parallel shell if the program you are shelling out
251 def __init__(self, dry = False, max = 10):
252 super(ParallelShell, self).__init__(dry=dry)
254 self.max = max # maximum of commands to run in parallel
256 def make(no_parallelize, max):
257 """Convenience method oriented towards command modules."""
259 return DummyParallelShell()
261 return ParallelShell(max=max)
262 def _async(self, proc, args, python, on_success, on_error, **kwargs):
264 Gets handed a :class:`subprocess.Proc` object from our deferred
265 execution. See :meth:`Shell.call` source code for details.
267 self.running[proc.pid] = (proc, args, python, on_success, on_error)
270 Blocking call that waits for an open subprocess slot. This is
271 automatically called by :meth:`Shell.call`.
273 # XXX: This API sucks; the actuall call/callAsUser call should
274 # probably block automatically (unless I have a good reason not to)
275 # bail out immediately on initial ramp up
276 if len(self.running) < self.max: return
277 # now, wait for open pids.
279 self.reap(*os.waitpid(-1, 0))
281 if e.errno == errno.ECHILD: return
284 """Waits for all of our subprocesses to terminate."""
287 self.reap(*os.waitpid(-1, 0))
289 if e.errno == errno.ECHILD: return
291 def reap(self, pid, status):
292 """Reaps a process."""
293 # ooh, zombie process. reap it
294 proc, args, python, on_success, on_error = self.running.pop(pid)
295 # XXX: this is slightly dangerous; should actually use
297 stdout = proc.stdout.read()
298 stderr = proc.stderr.read()
299 self._log(stdout, stderr)
301 if python: eclass = PythonCallError
302 else: eclass = CallError
303 on_error(eclass(proc.returncode, args, stdout, stderr))
305 on_success(stdout, stderr)
308 class DummyParallelShell(ParallelShell):
309 """Same API as :class:`ParallelShell`, but doesn't actually
310 parallelize (i.e. all calls to :meth:`wait` block.)"""
311 def __init__(self, dry = False):
312 super(DummyParallelShell, self).__init__(dry=dry, max=1)
314 class Error(wizard.Error):
315 """Base exception for this module"""
318 class CallError(Error):
319 """Indicates that a subprocess call returned a nonzero exit status."""
320 #: The exit code of the failed subprocess.
322 #: List of the program and arguments that failed.
324 #: The stdout of the program.
326 #: The stderr of the program.
328 def __init__(self, code, args, stdout, stderr):
334 compact = self.stderr.rstrip().split("\n")[-1]
335 return "%s (exited with %d)\n%s" % (compact, self.code, self.stderr)
337 class PythonCallError(CallError):
339 Indicates that a Python subprocess call had an uncaught exception.
340 This exception also contains the attributes of :class:`CallError`.
342 #: Name of the uncaught exception.
344 def __init__(self, code, args, stdout, stderr):
345 if stderr: self.name = util.get_exception_name(stderr)
346 CallError.__init__(self, code, args, stdout, stderr)
349 return "PythonCallError [%s]\n%s" % (self.name, self.stderr)
351 return "PythonCallError\n%s" % self.stderr