class Proc(object): def __init__(self, *args, **kwargs): self.timeout = kwargs.get('timeout', 0.1) self.display = kwargs.get('display', False) if 'host' in kwargs and 'port' in kwargs: self.p = socket.create_connection((kwargs['host'], kwargs['port'])) self.p.setblocking(0) else: self.p = Popen(args, stdin=PIPE, stdout=PIPE) fd = self.p.stdout.fileno() fl = fcntl.fcntl(fd, fcntl.F_GETFL) fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK) if kwargs.get('debug', False): raw_input("\x1b[32mpid %d is running, attach the debugger if needed. Hit enter key to continue...\x1b[0m" % self.p.pid) def write(self, s): if isinstance(self.p, Popen): select.select([], [self.p.stdin], []) self.p.stdin.write(s) else: select.select([], [self.p], []) self.p.sendall(s) if self.display: printable = re.sub(r'[^\s\x20-\x7e]', '.', s) sys.stdout.write("\x1b[33m%s\x1b[0m" % printable) # yellow sys.stdout.flush() def read(self, size=-1, timeout=None): if size < 0: chunk_size = 8192 buf = '' while True: chunk = self.read(chunk_size, timeout) buf += chunk if len(chunk) < chunk_size: break return buf if timeout is None: timeout = self.timeout if isinstance(self.p, Popen): stdout, read = self.p.stdout, self.p.stdout.read else: stdout, read = self.p, self.p.recv buf = '' while len(buf) < size: rlist, wlist, xlist = select.select([stdout], [], [], timeout) if not rlist: break chunk = read(size-len(buf)) if not chunk: break buf += chunk if self.display: printable = re.sub(r'[^\s\x20-\x7e]', '.', buf) sys.stdout.write("\x1b[36m%s\x1b[0m" % printable) # cyan sys.stdout.flush() return buf def read_until(self, s): buf = self.read(len(s), 864000) while not buf.endswith(s): buf += self.read(1, 864000) return buf def expect(self, regexp): buf = '' while not re.search(regexp, buf): buf += self.read(1, 864000) return buf def readline(self): return self.read_until('\n') def writeline(self, s): return self.write(s+'\n') def shutdown(self, writeonly=False): if isinstance(self.p, Popen): self.p.stdin.close() if not writeonly: self.p.stdout.close() else: if writeonly: self.p.shutdown(socket.SHUT_WR) else: self.p.shutdown(socket.SHUT_RDWR) def close(self): if isinstance(self.p, Popen): self.p.terminate() else: self.p.close() def wait(self, redirect_fd=None): check_cmd = 'echo "\x1b[32mgot a shell!\x1b[0m"' # green buf = self.read() sys.stdout.write(buf) if isinstance(self.p, Popen): if redirect_fd is not None: self.write(check_cmd + '\n') sys.stdout.write(self.read()) self.write('exec /bin/sh <&2 >&2\n') self.p.wait() return self.p.returncode else: if redirect_fd is not None: self.write(check_cmd + '\n') sys.stdout.write(self.read()) self.write("exec /bin/sh <&%(fd)d >&%(fd)d 2>&%(fd)d\n" % {'fd': redirect_fd}) t = Telnet() t.sock = self.p t.interact() t.close() @contextmanager def listen(self, port=4444, echotest=False): check_cmd = 'echo "\x1b[32mgot a shell!\x1b[0m"' # green s = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0) s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) s.bind(('', port)) # the empty string represents INADDR_ANY s.listen(1) if isinstance(self.p, Popen): addrinfo = socket.getaddrinfo('localhost', port, socket.AF_INET, socket.SOCK_STREAM) host = addrinfo[0][4][0] else: host = self.p.getsockname()[0] yield (host, port) c, addr = s.accept() s.close() if echotest: c.sendall(check_cmd + '\n') sys.stdout.write(c.recv(8192)) t = Telnet() t.sock = c t.interact() t.close() self.close() def pipe_output(self, *args): if isinstance(self.p, Popen): p_stdout = self.p.stdout else: p_stdout = self.p.makefile() p = Popen(args, stdin=p_stdout, stdout=PIPE) stdout, stderr = p.communicate() return stdout def write_p64(self, s): return self.write(p64(s)) def write_p32(self, s): return self.write(p32(s)) def read_p64(self, timeout=None): return p64(self.read(8, timeout)) def read_p32(self, timeout=None): return p32(self.read(4, timeout))