From 534002288c798a49c058e264cd497ffaaf05de48 Mon Sep 17 00:00:00 2001
From: Mitul16 <mitul.varshney.2@gmail.com>
Date: Sun, 18 Jul 2021 02:21:16 +0530
Subject: [PATCH] Added warning messages for 'KeyboardInterrupt'

I have added warning messages so that the attacker can use 'C-c' to cancel the session when used from command line (not from local prompt). And added the message for already handled same exception. I have also changed two warning message color from 'red' to 'yellow', because 'yellow' is being used everywhere in the code.
---
 pwncat/__main__.py          | 2 ++
 pwncat/commands/__init__.py | 4 ++--
 pwncat/commands/connect.py  | 2 ++
 3 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/pwncat/__main__.py b/pwncat/__main__.py
index 53590ec..e81c84d 100644
--- a/pwncat/__main__.py
+++ b/pwncat/__main__.py
@@ -277,6 +277,8 @@ def main():
                     )
                 except (ChannelError, PlatformError) as exc:
                     manager.log(f"connection failed: {exc}")
+                except KeyboardInterrupt:
+                    manager.log("[yellow]warning[/yellow]: cancelled by user")
 
         manager.interactive()
 
diff --git a/pwncat/commands/__init__.py b/pwncat/commands/__init__.py
index 8e91add..8dbf0e9 100644
--- a/pwncat/commands/__init__.py
+++ b/pwncat/commands/__init__.py
@@ -523,7 +523,7 @@ class CommandParser:
                 self.dispatch_line(command)
             except ChannelClosed as exc:
                 # A channel was unexpectedly closed
-                self.manager.log(f"[red]warning[/red]: {exc.channel}: channel closed")
+                self.manager.log(f"[yellow]warning[/yellow]: {exc.channel}: channel closed")
                 # Ensure any existing sessions are cleaned from the manager
                 exc.cleanup(self.manager)
             except pwncat.manager.InteractiveExit:
@@ -593,7 +593,7 @@ class CommandParser:
                 continue
             except ChannelClosed as exc:
                 # A channel was unexpectedly closed
-                self.manager.log(f"[red]warning[/red]: {exc.channel}: channel closed")
+                self.manager.log(f"[yellow]warning[/yellow]: {exc.channel}: channel closed")
                 # Ensure any existing sessions are cleaned from the manager
                 exc.cleanup(self.manager)
             except pwncat.manager.InteractiveExit:
diff --git a/pwncat/commands/connect.py b/pwncat/commands/connect.py
index aa7d3f1..f5648a0 100644
--- a/pwncat/commands/connect.py
+++ b/pwncat/commands/connect.py
@@ -244,3 +244,5 @@ class Command(CommandDefinition):
                 )
             except (ChannelError, PlatformError) as exc:
                 manager.log(f"connection failed: {exc}")
+            except KeyboardInterrupt:
+                manager.log("[yellow]warning[/yellow]: cancelled by user")