From aa21df29e76e73fcd5e5e16ba843413c54140950 Mon Sep 17 00:00:00 2001
From: Joe Testa <jtesta@positronsecurity.com>
Date: Mon, 24 May 2021 19:50:25 -0400
Subject: [PATCH] Now handles exceptions during server KEX parsing more
 gracefully.

---
 README.md                    |  1 +
 src/ssh_audit/gextest.py     | 12 +++++++++---
 src/ssh_audit/hostkeytest.py | 13 +++++++++----
 src/ssh_audit/ssh_audit.py   |  7 ++++++-
 4 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index f39b32a..be40370 100644
--- a/README.md
+++ b/README.md
@@ -177,6 +177,7 @@ For convenience, a web front-end on top of the command-line tool is available at
 ## ChangeLog
 ### v2.5.0-dev (???)
  - Fixed crash when running host key tests.
+ - Handles server connection failures more gracefully.
  - Now prints JSON with indents when `-jj` is used (useful for debugging).
  - Added MD5 fingerprints to verbose output.
  - Added `-d`/`--debug` option for getting debugging output; credit [Adam Russell](https://github.com/thecliguy).
diff --git a/src/ssh_audit/gextest.py b/src/ssh_audit/gextest.py
index 312e6aa..47efbce 100644
--- a/src/ssh_audit/gextest.py
+++ b/src/ssh_audit/gextest.py
@@ -26,6 +26,8 @@
 from typing import Dict, List, Set, Sequence, Tuple, Iterable  # noqa: F401
 from typing import Callable, Optional, Union, Any  # noqa: F401
 
+import traceback
+
 from ssh_audit.kexdh import KexGroupExchange_SHA1, KexGroupExchange_SHA256
 from ssh_audit.ssh2_kexdb import SSH2_KexDB
 from ssh_audit.ssh2_kex import SSH2_Kex
@@ -58,9 +60,13 @@ class GEXTest:
         # server's own values.
         s.send_kexinit(key_exchanges=[gex_alg], hostkeys=kex.key_algorithms, ciphers=kex.server.encryption, macs=kex.server.mac, compressions=kex.server.compression, languages=kex.server.languages)
 
-        # Parse the server's KEX.
-        _, payload = s.read_packet(2)
-        SSH2_Kex.parse(payload)
+        try:
+            # Parse the server's KEX.
+            _, payload = s.read_packet(2)
+            SSH2_Kex.parse(payload)
+        except Exception:
+            out.v("Failed to parse server's kex.  Stack trace:\n%s" % str(traceback.format_exc()), write_now=True)
+            return False
 
         return True
 
diff --git a/src/ssh_audit/hostkeytest.py b/src/ssh_audit/hostkeytest.py
index 8b0a6eb..6f87ac5 100644
--- a/src/ssh_audit/hostkeytest.py
+++ b/src/ssh_audit/hostkeytest.py
@@ -26,6 +26,8 @@
 from typing import Dict, List, Set, Sequence, Tuple, Iterable  # noqa: F401
 from typing import Callable, Optional, Union, Any  # noqa: F401
 
+import traceback
+
 from ssh_audit.kexdh import KexDH, KexGroup1, KexGroup14_SHA1, KexGroup14_SHA256, KexCurve25519_SHA256, KexGroup16_SHA512, KexGroup18_SHA512, KexGroupExchange_SHA1, KexGroupExchange_SHA256, KexNISTP256, KexNISTP384, KexNISTP521
 from ssh_audit.ssh2_kex import SSH2_Kex
 from ssh_audit.ssh2_kexdb import SSH2_KexDB
@@ -123,10 +125,13 @@ class HostKeyTest:
                     # Send our KEX using the specified group-exchange and most of the server's own values.
                     s.send_kexinit(key_exchanges=[kex_str], hostkeys=[host_key_type], ciphers=server_kex.server.encryption, macs=server_kex.server.mac, compressions=server_kex.server.compression, languages=server_kex.server.languages)
 
-                    # Parse the server's KEX.
-                    _, payload = s.read_packet()
-                    SSH2_Kex.parse(payload)
-
+                    try:
+                        # Parse the server's KEX.
+                        _, payload = s.read_packet()
+                        SSH2_Kex.parse(payload)
+                    except Exception:
+                        out.v("Failed to parse server's kex.  Stack trace:\n%s" % str(traceback.format_exc()), write_now=True)
+                        return
 
                 # Do the initial DH exchange.  The server responds back
                 # with the host key and its length.  Bingo.  We also get back the host key fingerprint.
diff --git a/src/ssh_audit/ssh_audit.py b/src/ssh_audit/ssh_audit.py
index d083ee5..d9da4f4 100755
--- a/src/ssh_audit/ssh_audit.py
+++ b/src/ssh_audit/ssh_audit.py
@@ -895,7 +895,12 @@ def audit(out: OutputBuffer, aconf: AuditConf, sshv: Optional[int] = None, print
     if sshv == 1:
         program_retval = output(out, aconf, banner, header, pkm=SSH1_PublicKeyMessage.parse(payload))
     elif sshv == 2:
-        kex = SSH2_Kex.parse(payload)
+        try:
+            kex = SSH2_Kex.parse(payload)
+        except Exception:
+            out.fail("Failed to parse server's kex.  Stack trace:\n%s" % str(traceback.format_exc()))
+            return exitcodes.CONNECTION_ERROR
+
         if aconf.client_audit is False:
             HostKeyTest.run(out, s, kex)
             GEXTest.run(out, s, kex)