Fixed exit command

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

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:

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

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.