diff --git a/pupy/modules/interactive_shell.py b/pupy/modules/interactive_shell.py index 9dfa5e14..7e9162ce 100644 --- a/pupy/modules/interactive_shell.py +++ b/pupy/modules/interactive_shell.py @@ -3,7 +3,6 @@ # Pupy is under the BSD 3-Clause license. see the LICENSE file at the root of the project for the detailed licence terms from pupylib.PupyModule import * -from pupylib.utils.rpyc_utils import redirected_stdio from rpyc.core.async import AsyncResultTimeout import sys import os @@ -17,8 +16,51 @@ if sys.platform!="win32": import array import time import StringIO -from threading import Event, Thread +from threading import Event, Thread, Lock import rpyc +import cmd + +class CmdRepl(cmd.Cmd): + def __init__(self, write_cb, completion): + self._write_cb = write_cb + self._complete = completion + self._write_lock = Lock() + self.prompt = '\r' + cmd.Cmd.__init__(self) + + def _con_write(self, data): + if not self._complete.is_set(): + with self._write_lock: + self.stdout.write(data) + self.stdout.flush() + if '\n' in data: + self.prompt = data.rsplit('\n', 1)[-1] + else: + self.prompt += data + + def do_EOF(self, line): + return True + + def precmd(self, line): + if self._complete.is_set(): + return 'EOF' + else: + return line + + def postcmd(self, stop, line): + if stop or self._complete.is_set(): + return True + + def emptyline(self): + pass + + def default(self, line): + with self._write_lock: + self._write_cb(line + '\n') + self.prompt = '' + + def postloop(self): + self._complete.set() __class_name__="InteractiveShell" @config(cat="admin") @@ -27,6 +69,8 @@ class InteractiveShell(PupyModule): open an interactive command shell. tty are well handled for targets running *nix """ max_clients=1 + pipe = None + complete = Event() def __init__(self, *args, **kwargs): PupyModule.__init__(self,*args, **kwargs) @@ -42,9 +86,9 @@ class InteractiveShell(PupyModule): fcntl.ioctl(pty.STDOUT_FILENO, termios.TIOCGWINSZ, buf, True) self.set_pty_size(buf[0], buf[1], buf[2], buf[3]) - def _start_read_loop(self, write_cb, complete): + def _start_read_loop(self, write_cb): t = Thread( - target=self._read_loop, args=(write_cb, complete) + target=self._read_loop, args=(write_cb,) ) t.daemon = True t.start() @@ -60,98 +104,133 @@ class InteractiveShell(PupyModule): buf.append(os.read(fd, 1)) return b''.join(buf) - def _read_loop(self, write_cb, complete): + def _read_loop(self, write_cb): try: - self._read_loop_base(write_cb, complete) - except AsyncResultTimeout: + self._read_loop_base(write_cb) + except AsyncResultTimeout, ReferenceError: pass finally: sys.stdout.write('\r\n') - complete.set() + self.complete.set() - - def _read_loop_base(self, write_cb, complete): + def _read_loop_base(self, write_cb): lastbuf = b'' write_cb = rpyc.async(write_cb) - while not complete.is_set(): + while not self.complete.is_set(): r, _, x = select.select([sys.stdin], [], [sys.stdin], None) if x: break if r: - if not complete.is_set(): + if not self.complete.is_set(): buf = self._read_stdin_non_block() if lastbuf.startswith(b'\r'): - vbuf = lastbuf + buf - if vbuf.startswith(b'\r~'): - if len(vbuf) < 3: - lastbuf = vbuf + lastbuf += buf + if lastbuf.startswith(b'\r~'): + if len(lastbuf) < 3: continue - elif vbuf.startswith(b'\r~.'): + elif lastbuf.startswith(b'\r~.'): break - elif vbuf.startswith(b'\r~,'): + elif lastbuf.startswith(b'\r~,'): self.client.conn._conn.ping(timeout=1) - buf = buf[3:] + buf = lastbuf[3:] if not buf: continue write_cb(buf) lastbuf = buf - def _remote_read(self, data, complete): - if not complete.is_set(): + def _remote_read(self, data): + if not self.complete.is_set(): os.write(sys.stdout.fileno(), data) def run(self, args): - if self.client.is_windows() or args.pseudo_tty: - self.client.load_package("interactive_shell") - encoding=None - program="/bin/sh" - if self.client.is_android(): - program="/system/bin/sh" - elif self.client.is_windows(): - program="cmd.exe" - if args.program: - program=args.program - with redirected_stdio(self.client.conn): - self.client.conn.modules.interactive_shell.interactive_open(program=program) - else: #handling tty - self.client.load_package("ptyshell") - - fd = sys.stdin.fileno() - old_settings = termios.tcgetattr(fd) - - ps = self.client.conn.modules['ptyshell'].PtyShell() - program = None - - if args.program: - program=args.program.split() - + if 'linux' in sys.platform and not args.pseudo_tty: try: - term = os.environ.get('TERM', 'xterm') - ps.spawn(program, term=term) - - closed = Event() - - self.set_pty_size=rpyc.async(ps.set_pty_size) - old_handler = pupylib.PupySignalHandler.set_signal_winch(self._signal_winch) - self._signal_winch(None, None) # set the remote tty sie to the current terminal size + fd = sys.stdin.fileno() + old_settings = termios.tcgetattr(fd) tty.setraw(fd) - - ps.start_read_loop(lambda data: self._remote_read(data, closed), closed.set) - self._start_read_loop(ps.write, closed) - - closed.wait() - - # Read loop here - - termios.tcsetattr(fd, termios.TCSADRAIN, old_settings) - pupylib.PupySignalHandler.set_signal_winch(old_handler) - + self.raw_pty(args) finally: - try: - self.ps.close() - except Exception: - pass - self.set_pty_size=None + termios.tcsetattr(fd, termios.TCSADRAIN, old_settings) + else: + # Well, this probably doesn't work at all + self.repl(args) + + def repl(self, args): + self.client.load_package('pupyutils.safepopen') + encoding=None + program="/bin/sh" + if self.client.is_android(): + program="/system/bin/sh" + elif self.client.is_windows(): + program="cmd.exe" + if args.program: + program=args.program + + self.pipe = self.client.conn.modules['pupyutils.safepopen'].SafePopen( + [ program ], + interactive=True, + ) + + print "DEBUG: {}".format(self.complete.is_set()) + + sys.stdout.write('\r\nREPL started. Ctrl-C will the module \r\n') + + repl = CmdRepl(self.pipe.write, self.complete) + self.pipe.execute(self.complete.set, repl._con_write) + + repl_thread = Thread(target=repl.cmdloop) + repl_thread.daemon = True + repl_thread.start() + + self.complete.wait() + self.pipe.terminate() + + # Well, there is no way to break upper thread without + # new 100500 threads which will wrap stdin, poll each other... + # Just press the fucked enter to avoid this crap + + sys.stdout.write('\r\nPress Enter to close to REPL\r\n') + + def raw_pty(self, args): + if self.client.is_windows(): + self.client.load_dll('winpty.dll') + self.client.load_package('winpty') + + self.client.load_package("ptyshell") + + ps = self.client.conn.modules['ptyshell'].PtyShell() + program = None + + if args.program: + program=args.program.split() + + try: + term = os.environ.get('TERM', 'xterm') + + ps.spawn(program, term=term) + + self.set_pty_size=rpyc.async(ps.set_pty_size) + old_handler = pupylib.PupySignalHandler.set_signal_winch(self._signal_winch) + self._signal_winch(None, None) # set the remote tty sie to the current terminal size + + self.complete = Event() + ps.start_read_loop(self._remote_read, self.complete.set) + self._start_read_loop(ps.write) + + self._signal_winch(None, None) + + self.complete.wait() + + finally: + pupylib.PupySignalHandler.set_signal_winch(old_handler) + try: + self.ps.close() + except Exception: + pass + self.set_pty_size=None + + def interrupt(self): + self.complete.set() diff --git a/pupy/packages/windows/all/ptyshell.py b/pupy/packages/windows/all/ptyshell.py new file mode 100644 index 00000000..5bb8dcf0 --- /dev/null +++ b/pupy/packages/windows/all/ptyshell.py @@ -0,0 +1,68 @@ +# -*- coding: utf-8 -*- + +import rpyc +import winpty +import threading + +class PtyShell(object): + def __init__(self): + self.pty = None + + def close(self): + if self.pty: + self.pty.close() + + def __del__(self): + self.close() + + def spawn(self, argv=None, term=None): + if self.pty: + return + + if not argv: + argv = r'C:\windows\system32\cmd.exe' + + self.pty = winpty.WinPTY(argv) + + def write(self, data): + if not self.pty: + return + + self.pty.write(data) + + def set_pty_size(self, ws_row, ws_col, ws_xpixel, ws_ypixel): + if not self.pty: + return + + self.pty.resize(ws_row, ws_col) + + def start_read_loop(self, print_callback, close_callback): + if not self.pty: + return + + t=threading.Thread( + target=self._read_loop, + args=(print_callback, close_callback) + ) + + t.daemon=True + t.start() + + def _read_loop(self, print_callback, close_callback): + cb = rpyc.async(print_callback) + close_cb = rpyc.async(close_callback) + + while True: + data = self.pty.read() + if not data: + break + + cb(data) + + close_cb() + + def close(self): + if not self.pty: + return + + self.pty.close() diff --git a/pupy/packages/windows/all/winpty.py b/pupy/packages/windows/all/winpty.py new file mode 100644 index 00000000..7cd2f1df --- /dev/null +++ b/pupy/packages/windows/all/winpty.py @@ -0,0 +1,210 @@ +# -*- coding: utf-8 -*- + +from contextlib import contextmanager + +from ctypes import * +from ctypes.wintypes import * + +from win32file import CreateFile, ReadFile, WriteFile, CloseHandle +from win32file import GENERIC_READ, GENERIC_WRITE, OPEN_EXISTING + +from win32pipe import PeekNamedPipe + +import pupy + +WINPTY_ERROR_SUCCESS = 0 +WINPTY_ERROR_OUT_OF_MEMORY = 1 +WINPTY_ERROR_SPAWN_CREATE_PROCESS_FAILED = 2 +WINPTY_ERROR_LOST_CONNECTION = 3 +WINPTY_ERROR_AGENT_EXE_MISSING = 4 +WINPTY_ERROR_UNSPECIFIED = 5 +WINPTY_ERROR_AGENT_DIED = 6 +WINPTY_ERROR_AGENT_TIMEOUT = 7 +WINPTY_ERROR_AGENT_CREATION_FAILED = 8 + +WINPTY_FLAG_CONERR = 0x1 +WINPTY_FLAG_PLAIN_OUTPUT = 0x2 +WINPTY_FLAG_COLOR_ESCAPES = 0x4 +WINPTY_FLAG_ALLOW_CURPROC_DESKTOP_CREATION = 0x8 + +WINPTY_MOUSE_MODE_NONE = 0 +WINPTY_MOUSE_MODE_AUTO = 1 +WINPTY_MOUSE_MODE_FORCE = 2 + +WINPTY_SPAWN_FLAG_AUTO_SHUTDOWN = 1 +WINPTY_SPAWN_FLAG_EXIT_AFTER_SHUTDOWN = 2 + +DLLNAME = 'WINPTY.DLL' + +_functions = { + 'winpty_error_code': CFUNCTYPE(DWORD, c_void_p), + 'winpty_error_msg': CFUNCTYPE(LPCWSTR, c_void_p), + 'winpty_error_free': CFUNCTYPE(None, c_void_p), + 'winpty_config_new': CFUNCTYPE(c_void_p, c_ulonglong, c_void_p), + 'winpty_config_free': CFUNCTYPE(None, c_void_p), + 'winpty_config_set_initial_size': CFUNCTYPE(None, c_void_p, c_int, c_int), + 'winpty_config_set_mouse_mode': CFUNCTYPE(None, c_void_p, c_int), + 'winpty_config_set_agent_timeout': CFUNCTYPE(None, c_void_p, c_uint), + 'winpty_open': CFUNCTYPE(c_void_p, c_void_p, c_void_p), + 'winpty_free': CFUNCTYPE(None, c_void_p), + 'winpty_agent_process': CFUNCTYPE(HWND, c_void_p), + 'winpty_spawn_config_new': CFUNCTYPE( + c_void_p, c_ulonglong, LPCWSTR, LPCWSTR, LPCWSTR, LPCWSTR, c_void_p), + 'winpty_spawn_config_free': CFUNCTYPE(None, c_void_p), + 'winpty_spawn': CFUNCTYPE(c_int, c_void_p, c_void_p, c_void_p, c_void_p, c_void_p, c_void_p), + 'winpty_set_size': CFUNCTYPE(c_int, c_void_p, c_int, c_int, c_void_p), + 'winpty_conin_name': CFUNCTYPE(LPCWSTR, c_void_p), + 'winpty_conout_name': CFUNCTYPE(LPCWSTR, c_void_p), + 'winpty_conerr_name': CFUNCTYPE(LPCWSTR, c_void_p) +} + +for funcname, definition in _functions.iteritems(): + funcaddr = pupy.find_function_address(DLLNAME, funcname) + if not funcaddr: + raise ImportError("Couldn't find function {} at winpty.dll".format(funcname)) + globals()[funcname] = definition(funcaddr) + +class WinPTYException(Exception): + def __init__(self, code, message): + Exception.__init__(self, message) + self.code = code + +@contextmanager +def winpty_error(): + error = c_void_p(None) + try: + yield pointer(error) + code = winpty_error_code(error) + if code != WINPTY_ERROR_SUCCESS: + message = winpty_error_msg(error) + raise WinPTYException(code, message) + finally: + winpty_error_free(error) + +class WinPTY(object): + def __init__(self, program, + cmdline=None, cwd=None, env=None, + spawn_flags=WINPTY_SPAWN_FLAG_AUTO_SHUTDOWN|WINPTY_SPAWN_FLAG_EXIT_AFTER_SHUTDOWN, + pty_flags=0, pty_size=(80,25), pty_mouse=WINPTY_MOUSE_MODE_NONE): + + self._closed = False + + config = None + try: + with winpty_error() as error: + config = winpty_config_new(pty_flags, error) + + cols, rows = pty_size + if cols and rows: + winpty_config_set_initial_size(config, cols, rows) + winpty_config_set_mouse_mode(config, pty_mouse) + + with winpty_error() as error: + self._pty = winpty_open(config, error) + finally: + winpty_config_free(config) + + self._conin = winpty_conin_name(self._pty) + self._conout = winpty_conout_name(self._pty) + self._conerr = winpty_conerr_name(self._pty) + + try: + self._conin_pipe = CreateFile( + self._conin, + GENERIC_WRITE, + 0, None, + OPEN_EXISTING, + 0, None + ) + + self._conout_pipe = CreateFile( + self._conout, + GENERIC_READ, + 0, None, + OPEN_EXISTING, + 0, None + ) + + if self._conerr: + self._conerr_pipe = CreateFile( + self._conerr, + GENERIC_READ, + 0, None, + OPEN_EXISTING, + 0, None + ) + else: + self._conerr_pipe = None + + try: + spawn_ctx = None + process_handle = HANDLE() + thread_handle = HANDLE() + create_process_error = DWORD() + + with winpty_error() as error: + spawn_ctx = winpty_spawn_config_new( + spawn_flags, program, cmdline, cwd, env, error + ) + + with winpty_error() as error: + winpty_spawn( + self._pty, spawn_ctx, + pointer(process_handle), + pointer(thread_handle), + pointer(create_process_error), + error + ) + + finally: + winpty_spawn_config_free(spawn_ctx) + + except: + self.close() + raise + + def write(self, data): + if self._closed: + return False + + try: + WriteFile(self._conin_pipe, data) + return True + except: + return False + + def read(self, amount=8192): + if self._closed: + return False + + try: + error, data = ReadFile(self._conout_pipe, amount) + except: + data = None + + return data + + def resize(self, cols, rows): + if self._closed: + return False + + with winpty_error() as error: + winpty_set_size(self._pty, rows, cols, error) + + def close(self): + if self._closed: + return False + + self._closed = True + + CloseHandle(self._conin_pipe) + CloseHandle(self._conout_pipe) + if self._conerr_pipe: + CloseHandle(self._conerr_pipe) + winpty_free(self._pty) + + def __enter__(self): + return self + + def __exit__(self, *args): + self.close() diff --git a/pupy/packages/windows/amd64/winpty.dll b/pupy/packages/windows/amd64/winpty.dll new file mode 100644 index 00000000..404b3621 Binary files /dev/null and b/pupy/packages/windows/amd64/winpty.dll differ diff --git a/pupy/packages/windows/x86/winpty.dll b/pupy/packages/windows/x86/winpty.dll new file mode 100644 index 00000000..ed5fdadc Binary files /dev/null and b/pupy/packages/windows/x86/winpty.dll differ