From 2c9a1dbc715f6c1ebce0e10952a15b282510fb39 Mon Sep 17 00:00:00 2001
From: Caleb Stewart <caleb.stewart94@gmail.com>
Date: Sun, 13 Jun 2021 22:09:41 -0400
Subject: [PATCH] Initial implementation of ssl-wrapped socket

---
 pwncat/channel/__init__.py |  7 ++++++-
 pwncat/channel/bind.py     |  2 ++
 pwncat/channel/socket.py   | 11 ++++++++++-
 pwncat/channel/ssl_bind.py | 22 ++++++++++++++++++++++
 test.py                    |  9 ++++-----
 5 files changed, 44 insertions(+), 7 deletions(-)
 create mode 100644 pwncat/channel/ssl_bind.py

diff --git a/pwncat/channel/__init__.py b/pwncat/channel/__init__.py
index 5fcaab1..e8cfaf3 100644
--- a/pwncat/channel/__init__.py
+++ b/pwncat/channel/__init__.py
@@ -581,7 +581,10 @@ def create(protocol: Optional[str] = None, **kwargs) -> Channel:
                 or kwargs["host"] == "0.0.0.0"
                 or kwargs["host"] is None
             ):
-                protocols.append("bind")
+                if "certfile" in kwargs or "keyfile" in kwargs:
+                    protocols.append("ssl-bind")
+                else:
+                    protocols.append("bind")
             else:
                 protocols.append("connect")
     else:
@@ -600,8 +603,10 @@ from pwncat.channel.ssh import Ssh  # noqa: E402
 from pwncat.channel.bind import Bind  # noqa: E402
 from pwncat.channel.socket import Socket  # noqa: E402
 from pwncat.channel.connect import Connect  # noqa: E402
+from pwncat.channel.ssl_bind import SSLBind  # noqa: E402
 
 register("socket", Socket)
 register("bind", Bind)
 register("connect", Connect)
 register("ssh", Ssh)
+register("ssl-bind", SSLBind)
diff --git a/pwncat/channel/bind.py b/pwncat/channel/bind.py
index eaf7b18..1a6c678 100644
--- a/pwncat/channel/bind.py
+++ b/pwncat/channel/bind.py
@@ -51,6 +51,8 @@ class Bind(Socket):
                 self._socket_connected(client)
             except KeyboardInterrupt:
                 raise ChannelError(self, "listener aborted")
+            except socket.error as exc:
+                raise ChannelError(self, str(exc))
             finally:
                 self.server.close()
 
diff --git a/pwncat/channel/socket.py b/pwncat/channel/socket.py
index befdc67..5239bd1 100644
--- a/pwncat/channel/socket.py
+++ b/pwncat/channel/socket.py
@@ -15,6 +15,7 @@ utilize this class to instantiate a session via an established socket.
         manager.interactive()
 """
 import os
+import ssl
 import errno
 import fcntl
 import socket
@@ -91,11 +92,14 @@ class Socket(Channel):
             while written < len(data):
                 try:
                     written += self.client.send(data[written:])
-                except BlockingIOError:
+                except (BlockingIOError, ssl.SSLWantWriteError, ssl.SSLWantReadError):
                     pass
         except BrokenPipeError as exc:
             self._connected = False
             raise ChannelClosed(self) from exc
+        except (ssl.SSLEOFError, ssl.SSLSyscallError, ssl.SSLZeroReturnError):
+            self._connected = False
+            raise ChannelClosed(self) from exc
 
         return len(data)
 
@@ -124,6 +128,11 @@ class Socket(Channel):
         try:
             data = data + self.client.recv(count)
             return data
+        except ssl.SSLWantReadError:
+            return data
+        except (ssl.SSLEOFError, ssl.SSLSyscallError, ssl.SSLZeroReturnError):
+            self._connected = False
+            raise ChannelClosed(self) from exc
         except socket.error as exc:
             if exc.args[0] == errno.EAGAIN or exc.args[0] == errno.EWOULDBLOCK:
                 return data
diff --git a/pwncat/channel/ssl_bind.py b/pwncat/channel/ssl_bind.py
new file mode 100644
index 0000000..8704677
--- /dev/null
+++ b/pwncat/channel/ssl_bind.py
@@ -0,0 +1,22 @@
+#!/usr/bin/env python3
+import ssl
+
+from pwncat.channel import ChannelError
+from pwncat.channel.bind import Bind
+
+
+class SSLBind(Bind):
+    def __init__(self, certfile: str = None, keyfile: str = None, **kwargs):
+        super().__init__(**kwargs)
+
+        self.context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+        self.context.load_cert_chain(certfile, keyfile)
+
+        self.server = self.context.wrap_socket(self.server)
+
+    def connect(self):
+
+        try:
+            super().connect()
+        except ssl.SSLError as exc:
+            raise ChannelError(self, str(exc))
diff --git a/test.py b/test.py
index 17b8b7a..ffb073c 100755
--- a/test.py
+++ b/test.py
@@ -19,12 +19,11 @@ with pwncat.manager.Manager("data/pwncatrc") as manager:
     # session = manager.create_session("windows", host="192.168.56.10", port=4444)
     # session = manager.create_session("windows", host="192.168.122.11", port=4444)
     # session = manager.create_session("linux", host="pwncat-ubuntu", port=4444)
-    session = manager.create_session("linux", host="127.0.0.1", port=4444)
+    # session = manager.create_session("linux", host="127.0.0.1", port=4444)
+    session = manager.create_session(
+        "linux", certfile="/tmp/cert.pem", keyfile="/tmp/cert.pem", port=4444
+    )
 
     # session.platform.powershell("amsiutils")
 
-    with open("/tmp/random", "rb") as source:
-        with session.platform.open("/tmp/random", "wb") as destination:
-            shutil.copyfileobj(source, destination)
-
     manager.interactive()