sploit: Add function Comm.readall_nonblock()

Function should consume all available incoming data from target and
return it, however will return 'immediately' (according to a
configurable timeout) if the pipe is empty.

Signed-off-by: Malfurious <m@lfurio.us>
Signed-off-by: dusoleil <howcansocksbereal@gmail.com>

1 files changed, 11 insertions, 0 deletions

diff --git a/sploit/comm.py b/sploit/comm.py
index 604045c..4b1d487 100644
--- a/sploit/comm.py
+++ b/sploit/comm.py
@@ -11,6 +11,7 @@ class Comm:
     logonread = True
     logonwrite = False
     flushonwrite = True
+    timeout = 0.25 # seconds
 
     def __init__(self, backend):
         self.back = backend
@@ -42,6 +43,16 @@ class Comm:
             pass
         return data
 
+    def readall_nonblock(self):
+        try:
+            os.set_blocking(self.back.stdin.fileno(), False)
+            poll = select.poll()
+            poll.register(self.back.stdin, select.POLLIN)
+            poll.poll(self.timeout)
+            return self.readall()
+        finally:
+            os.set_blocking(self.back.stdin.fileno(), True)
+
     def readuntil(self, pred, /, *args, **kwargs):
         data = b''
         pred = bind(pred, *args, **kwargs)