From e9ee01cd304eaed7c639ef36029d1298ffe4b54d Mon Sep 17 00:00:00 2001
From: Caleb Stewart <caleb.stewart94@gmail.com>
Date: Sun, 13 Jun 2021 02:32:37 -0400
Subject: [PATCH] Changed SSH Channel to use SSHClient interface

---
 pwncat/__main__.py    |  9 ++++-
 pwncat/channel/ssh.py | 82 +++++++++++++++++--------------------------
 2 files changed, 41 insertions(+), 50 deletions(-)

diff --git a/pwncat/__main__.py b/pwncat/__main__.py
index 954385e..75e351b 100644
--- a/pwncat/__main__.py
+++ b/pwncat/__main__.py
@@ -22,6 +22,13 @@ from pwncat.commands import connect
 from pwncat.platform import PlatformError
 
 
+def existing_file(path):
+    path = os.path.expanduser(path)
+    if os.path.isfile(path):
+        return path
+    raise NotADirectoryError(path)
+
+
 def main():
 
     # Default log-level is "INFO"
@@ -45,7 +52,7 @@ def main():
     parser.add_argument(
         "--identity",
         "-i",
-        type=argparse.FileType("r"),
+        type=existing_file,
         default=None,
         help="Private key for SSH authentication",
     )
diff --git a/pwncat/channel/ssh.py b/pwncat/channel/ssh.py
index dca02f2..827bb3c 100644
--- a/pwncat/channel/ssh.py
+++ b/pwncat/channel/ssh.py
@@ -40,60 +40,44 @@ class Ssh(Channel):
             password = prompt("Password: ", is_password=True)
 
         try:
-            # Connect to the remote host's ssh server
-            sock = socket.create_connection((host, port))
-        except Exception as exc:
-            raise ChannelError(str(exc))
+            client = paramiko.client.SSHClient()
+            client.set_missing_host_key_policy(paramiko.client.AutoAddPolicy)
 
-        # Create a paramiko SSH transport layer around the socket
-        t = paramiko.Transport(sock)
-        try:
-            t.start_client()
-        except paramiko.SSHException:
-            sock.close()
-            raise ChannelError("ssh negotiation failed")
-
-        if identity is not None:
             try:
-                # Load the private key for the user
-                if isinstance(identity, str):
-                    key = paramiko.RSAKey.from_private_key_file(
-                        os.path.expanduser(identity)
-                    )
-                else:
-                    key = paramiko.RSAKey.from_private_key(identity)
-            except:
-                password = prompt("RSA Private Key Passphrase: ", is_password=True)
-                try:
-                    key = paramiko.RSAKey.from_private_key_file(identity, password)
-                except:
-                    raise ChannelError("invalid private key or passphrase")
+                client.connect(
+                    hostname=host,
+                    port=port,
+                    username=user,
+                    password=password,
+                    key_filename=os.path.expanduser(identity),
+                    allow_agent=True,
+                    look_for_keys=False,
+                )
+            except paramiko.ssh_exception.PasswordRequiredException:
+                passphrase = prompt("RSA Private Key Passphrase: ", is_password=True)
+                client.connect(
+                    hostname=host,
+                    port=port,
+                    username=user,
+                    password=password,
+                    key_filename=os.path.expanduser(identity),
+                    allow_agent=True,
+                    look_for_keys=False,
+                    passphrase=passphrase,
+                )
 
-            # Attempt authentication
-            try:
-                t.auth_publickey(user, key)
-            except paramiko.ssh_exception.AuthenticationException as exc:
-                raise ChannelError(str(exc))
-        else:
-            try:
-                t.auth_password(user, password)
-            except paramiko.ssh_exception.AuthenticationException as exc:
-                raise ChannelError(str(exc))
+            columns, rows = os.get_terminal_size(0)
+            shell = client.invoke_shell(width=columns, height=rows)
+            shell.setblocking(0)
 
-        if not t.is_authenticated():
-            t.close()
-            sock.close()
-            raise ChannelError("authentication failed")
+            self.client = shell
+            self.address = (host, port)
+            self._connected = True
 
-        # Open an interactive session
-        chan = t.open_session()
-        chan.get_pty()
-        chan.invoke_shell()
-        chan.setblocking(0)
-
-        self.client = chan
-        self.address = (host, port)
-        self._connected = True
+        except paramiko.ssh_exception.AuthenticationException as exc:
+            raise ChannelError(f"ssh authentication failed: {str(exc)}") from exc
+        except (paramiko.ssh_exception.SSHException, socket.error) as exc:
+            raise ChannelError(f"ssh connection failed: {str(exc)}") from exc
 
     @property
     def connected(self):