From 3c33d015e8e5907931b7543fc04e538cdd261567 Mon Sep 17 00:00:00 2001
From: Caleb Stewart <caleb.stewart94@gmail.com>
Date: Thu, 17 Jun 2021 17:00:24 -0400
Subject: [PATCH] Added way for output thread to communicate channel error

---
 pwncat/manager.py           | 21 ++++++++++++++++++---
 pwncat/platform/__init__.py |  5 +++++
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/pwncat/manager.py b/pwncat/manager.py
index e50d7ea..4ebb254 100644
--- a/pwncat/manager.py
+++ b/pwncat/manager.py
@@ -16,6 +16,8 @@ even if there was an uncaught exception. The normal method of creating a manager
 """
 import os
 import sys
+import queue
+import signal
 import fnmatch
 import pkgutil
 import threading
@@ -576,8 +578,9 @@ class Manager:
 
             interactive_complete = threading.Event()
             output_thread = None
+            channel_error = None
 
-            def output_thread_main(target: Session):
+            def output_thread_main(target: Session, exception_queue: queue.SimpleQueue):
 
                 while not interactive_complete.is_set():
                     try:
@@ -593,14 +596,19 @@ class Manager:
                         else:
                             interactive_complete.wait(timeout=0.1)
 
-                    except ChannelError:
+                    except ChannelError as exc:
+                        exception_queue.put(exc)
                         interactive_complete.set()
+                        # This is a hack to get the interactive loop out of a blocking
+                        # read call. The interactive loop will receive a KeyboardInterrupt
+                        os.kill(os.getpid(), signal.SIGINT)
 
             try:
                 self.target.platform.interactive = True
 
+                exception_queue = queue.Queue(maxsize=1)
                 output_thread = threading.Thread(
-                    target=output_thread_main, args=[self.target]
+                    target=output_thread_main, args=[self.target, exception_queue]
                 )
                 output_thread.start()
 
@@ -608,6 +616,12 @@ class Manager:
                     self.target.platform.interactive_loop(interactive_complete)
                 except RawModeExit:
                     pass
+
+                try:
+                    raise exception_queue.get(block=False)
+                except queue.Empty:
+                    pass
+
                 self.target.platform.interactive = False
             except ChannelClosed:
                 self.log(
@@ -620,6 +634,7 @@ class Manager:
                 interactive_complete.set()
                 if output_thread is not None:
                     output_thread.join()
+                    output_thread.join()
 
     def create_session(self, platform: str, channel: Channel = None, **kwargs):
         """
diff --git a/pwncat/platform/__init__.py b/pwncat/platform/__init__.py
index dc703b4..74b7a5e 100644
--- a/pwncat/platform/__init__.py
+++ b/pwncat/platform/__init__.py
@@ -549,6 +549,11 @@ class Platform(ABC):
                 data = sys.stdin.buffer.read(64)
                 has_prefix = self.session.manager._process_input(data, has_prefix)
 
+        except KeyboardInterrupt:
+            # This is a hack to allow the output thread to signal completion
+            # We are in raw mode so this couldn't have come from a user pressing
+            # C-d.
+            pass
         finally:
             pwncat.util.pop_term_state()
             sys.stdin.reconfigure(line_buffering=False)