From 67cdcb2972be849fad777e208ee686604a3981ab Mon Sep 17 00:00:00 2001
From: Caleb Stewart <caleb.stewart94@gmail.com>
Date: Fri, 13 Nov 2020 12:16:22 -0500
Subject: [PATCH] Fixed exit command

---
 pwncat/__main__.py          | 2 +-
 pwncat/commands/__init__.py | 7 +++----
 pwncat/commands/exit.py     | 3 ++-
 pwncat/util.py              | 4 ++++
 4 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/pwncat/__main__.py b/pwncat/__main__.py
index 4dbd4c7..8f58e28 100644
--- a/pwncat/__main__.py
+++ b/pwncat/__main__.py
@@ -106,7 +106,7 @@ def main():
         except ConnectionResetError:
             pwncat.victim.restore_local_term()
             console.log("[yellow]warning[/yellow]: connection reset by remote host")
-        except SystemExit:
+        except (pwncat.util.CommandSystemExit, SystemExit):
             console.log("closing connection")
         finally:
             # Restore the shell
diff --git a/pwncat/commands/__init__.py b/pwncat/commands/__init__.py
index 806624a..2675f2e 100644
--- a/pwncat/commands/__init__.py
+++ b/pwncat/commands/__init__.py
@@ -210,10 +210,7 @@ class CommandParser:
 
                 if pwncat.config.module:
                     self.prompt.message = [
-                        (
-                            "fg:ansiyellow bold",
-                            f"({pwncat.config.module.name}) ",
-                        ),
+                        ("fg:ansiyellow bold", f"({pwncat.config.module.name}) ",),
                         ("fg:ansimagenta bold", "pwncat"),
                         ("", "$ "),
                     ]
@@ -233,6 +230,8 @@ class CommandParser:
             # We used to catch only KeyboardException, but this prevents a
             # badly written command from completely killing our remote
             # connection.
+            except pwncat.util.CommandSystemExit:
+                raise
             except EOFError:
                 # We don't have a connection yet, just exit
                 if pwncat.victim is None or pwncat.victim.client is None:
diff --git a/pwncat/commands/exit.py b/pwncat/commands/exit.py
index 2b97405..4363e99 100644
--- a/pwncat/commands/exit.py
+++ b/pwncat/commands/exit.py
@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+import pwncat
 from pwncat.util import console
 from pwncat.commands.base import CommandDefinition, Complete, Parameter
 
@@ -27,4 +28,4 @@ class Command(CommandDefinition):
             return
 
         # Get outa here!
-        raise EOFError
+        raise pwncat.util.CommandSystemExit
diff --git a/pwncat/util.py b/pwncat/util.py
index 0434c47..ddb4019 100644
--- a/pwncat/util.py
+++ b/pwncat/util.py
@@ -66,6 +66,10 @@ class Init(Enum):
     SYSV = auto()
 
 
+class CommandSystemExit(Exception):
+    """ A command has requested that we exit pwncat (mostly used for exit command) """
+
+
 class CompilationError(Exception):
     """
     Indicates that compilation failed on either the local or remote host.