pylib/csc/backends/ipc.py

   1 """
   2 IPC Library Functions
   3
   4 This module contains very UNIX-specific code to allow interactive
   5 communication with another program. For CEO they are required to
   6 talk to kadmin because there is no Kerberos administration Python
   7 module. Real bindings to libkadm5 are doable and thus a TODO.
   8 """
   9 import os, pty, select
  10
  11
  12 class _pty_file(object):
  13     """
  14     A 'file'-like wrapper class for pseudoterminal file descriptors.
  15
  16     This wrapper is necessary because Python has a nasty habit of throwing
  17     OSError at pty EOF.
  18
  19     This class also implements timeouts for read operations which are handy
  20     for avoiding deadlock when both processes are blocked in a read().
  21
  22     See the Python documentation of the file class for explanation
  23     of the methods.
  24     """
  25     def __init__(self, fd):
  26         self.fd = fd
  27         self.buffer = ''
  28         self.closed = False
  29     def __repr__(self):
  30         status = 'open'
  31         if self.closed:
  32             status = 'closed'
  33         return "<" + status + " pty '" + os.ttyname(self.fd) + "'>"
  34     def read(self, size=-1, block=True, timeout=0.1):
  35         if self.closed: raise ValueError
  36         if size < 0:
  37             data = None
  38
  39             # read data, catching OSError as EOF
  40             try:
  41                 while data != '':
  42
  43                     # wait timeout for the pty to become ready, otherwise stop reading
  44                     if not block and len(select.select([self.fd], [], [], timeout)[0]) == 0:
  45                         break
  46
  47                     data = os.read(self.fd, 65536)
  48                     self.buffer += data
  49             except OSError:
  50                 pass
  51
  52             data = self.buffer
  53             self.buffer = ''
  54             return data
  55         else:
  56             if len(self.buffer) < size:
  57
  58                 # read data, catching OSError as EOF
  59                 try:
  60
  61                     # wait timeout for the pty to become ready, then read
  62                     if block or len(select.select([self.fd], [], [], timeout)[0]) != 0:
  63                         self.buffer += os.read(self.fd, size - len(self.buffer) )
  64
  65                 except OSError:
  66                     pass
  67
  68             data = self.buffer[:size]
  69             self.buffer = self.buffer[size:]
  70             return data
  71     def readline(self, size=-1, block=True, timeout=0.1):
  72         data = None
  73
  74         # read data, catching OSError as EOF
  75         try:
  76             while data != '' and self.buffer.find("\n") == -1 and (size < 0 or len(self.buffer) < size):
  77
  78                 # wait timeout for the pty to become ready, otherwise stop reading
  79                 if not block and len(select.select([self.fd], [], [], timeout)[0]) == 0:
  80                     break
  81
  82                 data = os.read(self.fd, 128)
  83                 self.buffer += data
  84         except OSError:
  85             pass
  86
  87         split_index = self.buffer.find("\n") + 1
  88         if split_index < 0:
  89             split_index = len(self.buffer)
  90         if size >= 0 and split_index > size:
  91             split_index = size
  92         line = self.buffer[:split_index]
  93         self.buffer = self.buffer[split_index:]
  94         return line
  95     def readlines(self, sizehint=None, timeout=0.1):
  96         lines = []
  97         line = None
  98         while True:
  99             line = self.readline(-1, False, timeout)
 100             if line == '': break
 101             lines.append(line)
 102         return lines
 103     def write(self, data):
 104         if self.closed: raise ValueError
 105         os.write(self.fd, data)
 106     def writelines(self, lines):
 107         for line in lines:
 108             self.write(line)
 109     def __iter__(self):
 110         return self
 111     def next(self):
 112         line = self.readline()
 113         if line == '':
 114             raise StopIteration
 115         return line
 116     def isatty(self):
 117         if self.closed: raise ValueError
 118         return os.isatty(self.fd)
 119     def fileno(self):
 120         if self.closed: raise ValueError
 121         return self.fd
 122     def flush(self):
 123         if self.closed: raise ValueError
 124         os.fsync(self.fd)
 125     def close(self):
 126         if not self.closed: os.close(self.fd)
 127         self.closed = True
 128
 129
 130 def popeni(command, args, env=None):
 131     """
 132     Open an interactive session with another command.
 133
 134     Parameters:
 135         command - the command to run (full path)
 136         args    - a list of arguments to pass to command
 137         env     - optional environment for command
 138
 139     Returns: (pid, stdout, stdin)
 140     """
 141
 142     # use a pipe to send data to the child
 143     child_stdin, parent_stdin = os.pipe()
 144
 145     # a pipe for receiving data would cause buffering and
 146     # is therefore not suitable for interactive communication
 147     # i.e. parent_stdout, child_stdout = os.pipe()
 148
 149     # therefore a pty must be used instead
 150     master, slave = pty.openpty()
 151
 152     # collect both stdout and stderr on the pty
 153     parent_stdout, child_stdout = master, slave
 154     parent_stderr, child_stderr = master, slave
 155
 156     # fork the child to communicate with
 157     pid = os.fork()
 158
 159     # child process
 160     if pid == 0:
 161
 162         # close all of the parent's fds
 163         os.close(parent_stdin)
 164         if parent_stdout != parent_stdin:
 165             os.close(parent_stdout)
 166         if parent_stderr != parent_stdin and parent_stderr != parent_stdout:
 167             os.close(parent_stderr)
 168
 169         # if stdout is a terminal, set it to the controlling terminal
 170         if os.isatty(child_stdout):
 171
 172             # determine the filename of the tty
 173             tty = os.ttyname(child_stdout)
 174
 175             # create a new session to disconnect
 176             # from the parent's controlling terminal
 177             os.setsid()
 178
 179             # set the controlling terminal to the pty
 180             # by opening it (and closing it again since
 181             # it's already open as child_stdout)
 182             fd = os.open(tty, os.O_RDWR)
 183             os.close(fd)
 184
 185         # init stdin/out/err
 186         os.dup2(child_stdin,  0)
 187         os.dup2(child_stdout, 1)
 188         if child_stderr >= 0:
 189             os.dup2(child_stderr, 2)
 190
 191         # finally, execute the child
 192         if env:
 193             os.execv(command, args, env)
 194         else:
 195             os.execv(command, args)
 196
 197     # parent process
 198     else:
 199
 200         # close all of the child's fds
 201         os.close(child_stdin)
 202         if child_stdout != child_stdin:
 203             os.close(child_stdout)
 204         if child_stderr >= 0 and child_stderr != child_stdin and child_stderr != child_stdout:
 205             os.close(child_stderr)
 206
 207         return pid, _pty_file(parent_stdout), os.fdopen(parent_stdin, 'w')
 208
 209
 210
 211 ### Tests ###
 212
 213 if __name__ == '__main__':
 214
 215     from csc.common.test import *
 216
 217     prog = '/bin/cat'
 218     argv = [ prog ]
 219     message = "test\n"
 220
 221     test(popeni)
 222     proc, recv, send = popeni(prog, argv)
 223     send.write(message)
 224     send.flush()
 225     line = recv.readline()
 226     assert_equal(message.strip(), line.strip())
 227     success()