From a95f2df50c3009f8584ceec6ebfaa93f55f22ed9 Mon Sep 17 00:00:00 2001
From: Caleb Stewart <caleb.stewart94@gmail.com>
Date: Fri, 18 Jun 2021 19:45:59 -0400
Subject: [PATCH] Fixed Windows platform transitions

Exception handling in the output thread was cleaned up and had Windows
platform raise the RawModeExit exception to trigger an exit when
interactive end marker was observed.
---
 pwncat/manager.py          | 12 ++++++------
 pwncat/platform/windows.py | 14 ++++++--------
 2 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/pwncat/manager.py b/pwncat/manager.py
index d01f75a..4354f9a 100644
--- a/pwncat/manager.py
+++ b/pwncat/manager.py
@@ -590,12 +590,9 @@ class Manager:
                             data = target.platform.channel.recv(4096)
 
                             if data != b"" and data is not None:
-                                try:
-                                    data = target.platform.process_output(data)
-                                    sys.stdout.buffer.write(data)
-                                    sys.stdout.buffer.flush()
-                                except RawModeExit:
-                                    interactive_complete.set()
+                                data = target.platform.process_output(data)
+                                sys.stdout.buffer.write(data)
+                                sys.stdout.buffer.flush()
                             else:
                                 interactive_complete.wait(timeout=0.1)
 
@@ -605,6 +602,9 @@ class Manager:
                             # 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)
+                        except RawModeExit:
+                            interactive_complete.set()
+                            os.kill(os.getpid(), signal.SIGINT)
 
                 try:
                     self.target.platform.interactive = True
diff --git a/pwncat/platform/windows.py b/pwncat/platform/windows.py
index 2221bf8..0e02c8b 100644
--- a/pwncat/platform/windows.py
+++ b/pwncat/platform/windows.py
@@ -13,7 +13,6 @@ processes and open multiple files with this platform. However, you should be
 careful to cleanup all processes and files prior to return from your method
 or code as the C2 will not attempt to garbage collect file or proces handles.
 """
-import os
 import sys
 import gzip
 import json
@@ -46,7 +45,7 @@ PWNCAT_WINDOWS_C2_VERSION = "v0.2.1"
 PWNCAT_WINDOWS_C2_RELEASE_URL = "https://github.com/calebstewart/pwncat-windows-c2/releases/download/{version}/pwncat-windows-{version}.tar.gz"
 
 
-class PowershellError(Exception):
+class PowershellError(PlatformError):
     """Executing a powershell script caused an error"""
 
     def __init__(self, msg):
@@ -55,7 +54,7 @@ class PowershellError(Exception):
         self.message = msg
 
 
-class ProtocolError(Exception):
+class ProtocolError(PlatformError):
     def __init__(self, code: int, message: str):
         self.code = code
         self.message = message
@@ -908,7 +907,7 @@ function prompt {
         transformed = bytearray(b"")
         has_cr = False
 
-        for b in data:
+        for idx, b in enumerate(data):
 
             # Basically, we just transform bare \r to \r\n
             if has_cr and b != ord("\n"):
@@ -924,10 +923,9 @@ function prompt {
             if INTERACTIVE_END_MARKER[self.interactive_tracker] == b:
                 self.interactive_tracker += 1
                 if self.interactive_tracker == len(INTERACTIVE_END_MARKER):
-                    # NOTE: this is a dirty hack to trigger the main input thread
-                    # to leave interactive mode, because it's bound in an input call
-                    os.kill(os.getpid(), signal.SIGINT)
-                    raise pwncat.manager.RawModeExit
+                    self.interactive_tracker = 0
+                    self.channel.unrecv(data[idx + 1 :])
+                    raise pwncat.util.RawModeExit
             else:
                 self.interactive_tracker = 0