add ptyexec helper script to run subprocesses in pseudo ttys and unbuffer

their I/O.



git-svn-id: http://munki.googlecode.com/svn/trunk@1284 a4e17f2e-e282-11dd-95e1-755cbddbdd66

code/client/ptyexec

@@ -0,0 +1,128 @@
+#!/usr/bin/python
+# encoding: utf-8
+#
+# Copyright 2011 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the 'License');
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an 'AS IS' BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+ptyexec
+
+Created by John Randolph 2011-08-11.
+
+Utility script to run a subprocess in a pseudo tty.
+
+This will have the effect of unbuffering output I/O from the subprocess.
+stdin of the subprocess is not connected to the stdin of this parent
+process.
+"""
+
+import fcntl
+import os
+import pty
+import select
+import signal
+import sys
+
+
+child_exited = {}
+
+
+def SetFileNonBlocking(f, non_blocking=True):
+    """Set non-blocking flag on a file object.
+
+    Args:
+        f: file
+        non_blocking: bool, default True, non-blocking mode or not
+    """
+    flags = fcntl.fcntl(f.fileno(), fcntl.F_GETFL)
+    if bool(flags & os.O_NONBLOCK) != non_blocking:
+        flags ^= os.O_NONBLOCK
+    fcntl.fcntl(f.fileno(), fcntl.F_SETFL, flags)
+
+
+def SigHandler(signum, frame):
+    """Handle a signal.
+
+    Args:
+        signum: int, signal number
+        frame: frame, stack frame where signal was received
+    """
+    global child_exited
+
+    if signum == signal.SIGCHLD:
+        x = os.waitpid(-1, 0)
+        if x[0] > 0:
+            child_exited[x[0]] = x[1] >> 8    # get exit status from LSB
+
+
+def Usage(arg0):
+    """Print usage."""
+    print >>sys.stderr, 'Usage: %s [command to run] [arguments...]' % arg0
+    return 0
+
+
+def PtyExec(argv):
+    """Setup pty and exec argv.
+
+    Args:
+        argv: list, arguments
+    Returns:
+        int, status code from child process upon completion, or 1 if a fork
+            error occurs.
+        never returns on the child side of the fork.
+    """
+    signal.signal(signal.SIGCHLD, SigHandler)
+    pid, fd = pty.fork()
+
+    if pid == 0:         # child
+        try:
+            os.execv(argv[0], argv)
+        except OSError, e:
+            print >>sys.stderr, str(e)
+            sys.exit(1)
+    elif pid > 0:        # parent
+        f = os.fdopen(fd, 'r+', 0)
+        SetFileNonBlocking(f)
+        while 1:
+            try:
+                (rl, wl, xl) = select.select([f], [], [], 5.0)
+            except select.error, e:
+                rl = []
+
+            if f in rl:
+                l = f.read()
+                sys.stdout.write(l)
+                sys.stdout.flush()
+
+            if pid in child_exited:
+                break
+
+        f.close()
+    elif pid == -1:    # error, never forked.
+        print >>sys.stderr, 'fork() error'
+        return 1
+
+    return child_exited.get(pid, 1)
+
+
+def main(argv):
+    """Main."""
+    if len(argv) < 2:
+        return Usage(argv[0])
+
+    argv.pop(0)
+    return PtyExec(argv)
+
+
+if __name__ == '__main__':
+    sys.exit(main(sys.argv))