From 38a8b3f528d9e02fd08d7d98a5c4beb530700cc2 Mon Sep 17 00:00:00 2001 From: Malfurious Date: Thu, 10 Mar 2022 05:23:00 -0500 Subject: sploit: Clean up function Comm.interact() The previous patches in this series have needed to utilize similar logic as Comm.interact() throughout other parts of the Comm class. This patch just revisits .interact() to clean up redundant code. Co-authored-by: dusoleil Signed-off-by: Malfurious Signed-off-by: dusoleil --- tools/sploit/sploit/comm.py | 58 ++++++++++++++++++++------------------------- 1 file changed, 26 insertions(+), 32 deletions(-) diff --git a/tools/sploit/sploit/comm.py b/tools/sploit/sploit/comm.py index 3972f16..265ab96 100644 --- a/tools/sploit/sploit/comm.py +++ b/tools/sploit/sploit/comm.py @@ -88,41 +88,35 @@ class Comm: self.write(data + b'\n') def interact(self): - ilog("<--Interact Mode-->") stdin = sys.stdin.buffer - os.set_blocking(self.back.stdin.fileno(), False) - os.set_blocking(stdin.fileno(), False) - poll = select.poll() - poll.register(self.back.stdin, select.POLLIN) - poll.register(stdin, select.POLLIN) - brk = False - def readall(read, write): - while(True): - data = read() - if(data == b''): - break - write(data) - def writeinput(write): - ilog(write, file=sys.stdout, color=NORMAL) + event = select.POLLIN + + def readall_stdin(): + for line in stdin: + self.write(line) + readtable = { - stdin.fileno() : lambda : readall(stdin.readline, self.write), - self.back.stdin.fileno() : lambda : readall(self.back.stdin.readline, writeinput) + self.back.stdin.fileno(): self.readall_nonblock, + stdin.fileno(): readall_stdin, } - readtable[self.back.stdin.fileno()]() - while(not brk): - try: - ioevents = poll.poll(100) - for ev in ioevents: - if(ev[1] & select.POLLIN): - readtable[ev[0]]() - else: - brk = True - break - except KeyboardInterrupt: - break - os.set_blocking(self.back.stdin.fileno(), True) - os.set_blocking(stdin.fileno(), True) - ilog("<--Interact Mode Done-->") + + try: + ilog("<--Interact Mode-->") + os.set_blocking(stdin.fileno(), False) + + poll = select.poll() + poll.register(self.back.stdin, event) + poll.register(stdin, event) + + while True: + for fd, e in poll.poll(self.timeout): + if not e & event: return + readtable[fd]() + except KeyboardInterrupt: + pass + finally: + os.set_blocking(stdin.fileno(), True) + ilog("<--Interact Mode Done-->") def popen(cmdline=''): io = Comm((Process(cmdline.split()) if len(cmdline) > 0 else Pipes())) -- cgit v1.2.3