From 02f96b7bd8b11357125821b53c1d9412804e2464 Mon Sep 17 00:00:00 2001 From: Wolfgang Hotwagner Date: Tue, 12 Mar 2024 11:39:56 +0100 Subject: [PATCH] use fctnl for non-blocking readline of subprocess --- .../executors/shell/shellexecutor.py | 25 ++++++++++++------- 1 file changed, 16 insertions(+), 9 deletions(-) diff --git a/src/attackmate/executors/shell/shellexecutor.py b/src/attackmate/executors/shell/shellexecutor.py index fae94df..5d9d5ff 100644 --- a/src/attackmate/executors/shell/shellexecutor.py +++ b/src/attackmate/executors/shell/shellexecutor.py @@ -5,6 +5,7 @@ commands in AttackMate. """ +import os import subprocess from subprocess import TimeoutExpired from datetime import datetime @@ -47,9 +48,18 @@ def popen_close(self, proc): proc.wait(timeout=10) @staticmethod - def enqueue_output(stdout, queue): - for line in iter(stdout.readline, b''): - queue.put(line) + def non_block_read(stdout): + fd = stdout.fileno() + try: + import fcntl + fl = fcntl.fcntl(fd, fcntl.F_GETFL) + fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK) + except ImportError: + raise ExecException("The 'fcntl' module is not available. This functionality requires a Unix-like operating system.") + try: + return stdout.read() + except: + return b"" def popen_noninteractive(self, proc: subprocess.Popen, cmd: bytes, timeout=None) -> str: self.logger.debug("Running non interactive command") @@ -71,16 +81,13 @@ def popen_interactive(self, proc: subprocess.Popen, cmd: bytes, timeout: int=5, proc.stdin.write(cmd) proc.stdin.flush() - t = Thread(target=ShellExecutor.enqueue_output, args=(proc.stdout, q)) - t.daemon = True - t.start() - outline = b'' if read: begin = datetime.now() while (datetime.now() - begin ).total_seconds() < timeout: - if q.qsize() > 0: - outline += q.get_nowait() + tmp = self.non_block_read(proc.stdout) + if tmp: + outline += tmp begin = datetime.now() # reset timer when data comes return outline.decode()