Advertisement
cgrunwald

Untitled

Oct 21st, 2010
126
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 3.56 KB | None | 0 0
  1. import os
  2. import subprocess
  3. import errno
  4. import time
  5. import sys
  6.  
  7. PIPE = subprocess.PIPE
  8.  
  9. if subprocess.mswindows:
  10.     from win32file import ReadFile, WriteFile
  11.     from win32pipe import PeekNamedPipe
  12.     import msvcrt
  13. else:
  14.     import select
  15.     import fcntl
  16.  
  17. class Popen(subprocess.Popen):
  18.     def recv(self, maxsize=None):
  19.         return self._recv('stdout', maxsize)
  20.  
  21.     def recv_err(self, maxsize=None):
  22.         return self._recv('stderr', maxsize)
  23.  
  24.     def send_recv(self, input='', maxsize=None):
  25.         return self.send(input), self.recv(maxsize), self.recv_err(maxsize)
  26.  
  27.     def get_conn_maxsize(self, which, maxsize):
  28.         if maxsize is None:
  29.             maxsize = 1024
  30.         elif maxsize < 1:
  31.             maxsize = 1
  32.         return getattr(self, which), maxsize
  33.  
  34.     def _close(self, which):
  35.         getattr(self, which).close()
  36.         setattr(self, which, None)
  37.  
  38.     if subprocess.mswindows:
  39.         def send(self, input):
  40.             if not self.stdin:
  41.                 return None
  42.  
  43.             try:
  44.                 x = msvcrt.get_osfhandle(self.stdin.fileno())
  45.                 (errCode, written) = WriteFile(x, input)
  46.             except ValueError:
  47.                 return self._close('stdin')
  48.             except (subprocess.pywintypes.error, Exception), why:
  49.                 if why[0] in (109, errno.ESHUTDOWN):
  50.                     return self._close('stdin')
  51.                 raise
  52.  
  53.             return written
  54.  
  55.         def _recv(self, which, maxsize):
  56.             conn, maxsize = self.get_conn_maxsize(which, maxsize)
  57.             if conn is None:
  58.                 return None
  59.  
  60.             try:
  61.                 x = msvcrt.get_osfhandle(conn.fileno())
  62.                 (read, nAvail, nMessage) = PeekNamedPipe(x, 0)
  63.                 if maxsize < nAvail:
  64.                     nAvail = maxsize
  65.                 if nAvail > 0:
  66.                     (errCode, read) = ReadFile(x, nAvail, None)
  67.             except ValueError:
  68.                 return self._close(which)
  69.             except (subprocess.pywintypes.error, Exception), why:
  70.                 if why[0] in (109, errno.ESHUTDOWN):
  71.                     return self._close(which)
  72.                 raise
  73.  
  74.             if self.universal_newlines:
  75.                 read = self._translate_newlines(read)
  76.             return read
  77.  
  78.     else:
  79.         def send(self, input):
  80.             if not self.stdin:
  81.                 return None
  82.  
  83.             if not select.select([], [self.stdin], [], 0)[1]:
  84.                 return 0
  85.  
  86.             try:
  87.                 written = os.write(self.stdin.fileno(), input)
  88.             except OSError, why:
  89.                 if why[0] == errno.EPIPE: #broken pipe
  90.                     return self._close('stdin')
  91.                 raise
  92.  
  93.             return written
  94.  
  95.         def _recv(self, which, maxsize):
  96.             conn, maxsize = self.get_conn_maxsize(which, maxsize)
  97.             if conn is None:
  98.                 return None
  99.  
  100.             flags = fcntl.fcntl(conn, fcntl.F_GETFL)
  101.             if not conn.closed:
  102.                 fcntl.fcntl(conn, fcntl.F_SETFL, flags| os.O_NONBLOCK)
  103.  
  104.             try:
  105.                 if not select.select([conn], [], [], 0)[0]:
  106.                     return ''
  107.  
  108.                 r = conn.read(maxsize)
  109.                 if not r:
  110.                     return self._close(which)
  111.  
  112.                 if self.universal_newlines:
  113.                     r = self._translate_newlines(r)
  114.                 return r
  115.             finally:
  116.                 if not conn.closed:
  117.                     fcntl.fcntl(conn, fcntl.F_SETFL, flags)
  118.  
  119. message = "Other end disconnected!"
  120.  
  121. def recv_some(p, t=.1, e=1, tr=5, stderr=0):
  122.     if tr < 1:
  123.         tr = 1
  124.     x = time.time()+t
  125.     y = []
  126.     r = ''
  127.     pr = p.recv
  128.     if stderr:
  129.         pr = p.recv_err
  130.     while time.time() < x or r:
  131.         r = pr()
  132.         if r is None:
  133.             if e:
  134.                 raise Exception(message)
  135.             else:
  136.                 break
  137.         elif r:
  138.             y.append(r)
  139.         else:
  140.             time.sleep(max((x-time.time())/tr, 0))
  141.     return ''.join(y)
  142.  
  143. def send_all(p, data):
  144.     while len(data):
  145.         sent = p.send(data)
  146.         if sent is None:
  147.             raise Exception(message)
  148.         data = buffer(data, sent)
  149.  
  150. if __name__ == '__main__':
  151.     if sys.platform == 'win32':
  152.         shell, commands, tail = ('cmd', ('dir /w', 'echo HELLO WORLD'), '\r\n')
  153.     else:
  154.         shell, commands, tail = ('sh', ('ls', 'echo HELLO WORLD'), '\n')
  155.  
  156.     a = Popen(shell, stdin=PIPE, stdout=PIPE)
  157.     print recv_some(a),
  158.     for cmd in commands:
  159.         send_all(a, cmd + tail)
  160.         print recv_some(a),
  161.     send_all(a, 'exit' + tail)
  162.     print recv_some(a, e=0)
  163.     a.wait()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement