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.