From d621880324f1f6fe41084253da2c34cc6a357b6e Mon Sep 17 00:00:00 2001
From: Caleb Stewart <caleb.stewart94@gmail.com>
Date: Sun, 26 Dec 2021 15:29:48 -0500
Subject: [PATCH] Fixed SSL argument parsing for bind channels

---
 CHANGELOG.md                  |  1 +
 pwncat/__main__.py            |  8 +++++++-
 pwncat/channel/ssl_bind.py    | 10 +++++++++-
 pwncat/channel/ssl_connect.py |  2 +-
 pwncat/commands/connect.py    | 12 ++++++++++--
 5 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 066d535..e824582 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -15,6 +15,7 @@ and simply didn't have the time to go back and retroactively create one.
 - Utilized Paramiko SSHClient which will also utilize the SSHAgent if available by default and supports key types aside from RSA ([#91](https://github.com/calebstewart/pwncat/issues/91))
 - Added implant module `list` command to match documentation ([#224](https://github.com/calebstewart/pwncat/issues/224)).
 - Update documentation to clarify implant reconnection
+- Fixed `--ssl` argument parsing for bind channels.
 
 ## [0.5.1] - 2021-12-07
 Minor bug fixes. Mainly typos from changing the package name.
diff --git a/pwncat/__main__.py b/pwncat/__main__.py
index 88152df..9736980 100644
--- a/pwncat/__main__.py
+++ b/pwncat/__main__.py
@@ -233,11 +233,17 @@ def main():
             if query_args["certfile"] is not None or query_args["keyfile"] is not None:
                 query_args["ssl"] = True
 
-            if query_args["protocol"] not in [None, "bind", "connect"] and args.ssl:
+            if query_args["protocol"] not in [
+                None,
+                "bind",
+                "connect",
+            ] and query_args.get("ssl"):
                 console.log(
                     f"[red]error[/red]: --ssl is incompatible with an [yellow]{query_args['protocol']}[/yellow] protocol"
                 )
                 return
+            elif query_args["protocol"] is not None:
+                query_args["protocol"] = "ssl-" + query_args["protocol"]
 
             if (
                 sum(
diff --git a/pwncat/channel/ssl_bind.py b/pwncat/channel/ssl_bind.py
index 9ffbdb8..170e547 100644
--- a/pwncat/channel/ssl_bind.py
+++ b/pwncat/channel/ssl_bind.py
@@ -23,7 +23,15 @@ class SSLBind(Bind):
 
         self.context.load_cert_chain(certfile, keyfile)
 
-        self.server = self.context.wrap_socket(self.server)
+        # self.server = self.context.wrap_socket(self.server)
+
+    def _socket_connected(self, client):
+        try:
+            client = self.context.wrap_socket(client, server_side=True)
+        except ssl.SSLError as exc:
+            raise ChannelError(self, str(exc))
+
+        super()._socket_connected(client)
 
     def connect(self):
 
diff --git a/pwncat/channel/ssl_connect.py b/pwncat/channel/ssl_connect.py
index 28d8c71..1a52cf9 100644
--- a/pwncat/channel/ssl_connect.py
+++ b/pwncat/channel/ssl_connect.py
@@ -17,6 +17,6 @@ class SSLConnect(Connect):
 
             client = self.context.wrap_socket(client)
         except ssl.SSLError as exc:
-            raise ChannelError(str(exc))
+            raise ChannelError(self, str(exc))
 
         super()._socket_connected(client)
diff --git a/pwncat/commands/connect.py b/pwncat/commands/connect.py
index 2c94c51..85ba17c 100644
--- a/pwncat/commands/connect.py
+++ b/pwncat/commands/connect.py
@@ -195,6 +195,16 @@ class Command(CommandDefinition):
         if query_args["certfile"] is not None or query_args["keyfile"] is not None:
             query_args["ssl"] = True
 
+        if query_args["protocol"] not in [None, "bind", "connect"] and query_args.get(
+            "ssl"
+        ):
+            console.log(
+                f"[red]error[/red]: --ssl is incompatible with an [yellow]{query_args['protocol']}[/yellow] protocol"
+            )
+            return
+        elif query_args["protocol"] is not None:
+            query_args["protocol"] = "ssl-" + query_args["protocol"]
+
         if (
             sum(
                 [
@@ -208,8 +218,6 @@ class Command(CommandDefinition):
             console.log("[red]error[/red]: multiple ports specified")
             return
 
-        console.log(args.pos_port)
-
         if args.port is not None:
             query_args["port"] = args.port
         if args.pos_port is not None: