From c483fe1861bcfaefabec21a9195b7c226540aaa4 Mon Sep 17 00:00:00 2001
From: Joe Testa <jtesta@positronsecurity.com>
Date: Fri, 26 Feb 2021 16:01:30 -0500
Subject: [PATCH] Fixed a crash while doing host key tests.

---
 README.md                    | 3 +++
 src/ssh_audit/globals.py     | 2 +-
 src/ssh_audit/hostkeytest.py | 2 +-
 src/ssh_audit/kexdh.py       | 2 +-
 4 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 780144b..a11ba14 100644
--- a/README.md
+++ b/README.md
@@ -161,6 +161,9 @@ $ docker pull positronsecurity/ssh-audit
 For convenience, a web front-end on top of the command-line tool is available at [https://www.ssh-audit.com/](https://www.ssh-audit.com/).
 
 ## ChangeLog
+### v2.5.0-dev (???)
+ - Fixed crash when running host key tests.
+
 ### v2.4.0 (2021-02-23)
  - Added multi-threaded scanning support.
  - Added built-in Windows manual page (see `-m`/`--manual`); credit [Adam Russell](https://github.com/thecliguy).
diff --git a/src/ssh_audit/globals.py b/src/ssh_audit/globals.py
index 8e8fdea..874d418 100644
--- a/src/ssh_audit/globals.py
+++ b/src/ssh_audit/globals.py
@@ -21,7 +21,7 @@
    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    THE SOFTWARE.
 """
-VERSION = 'v2.4.0'
+VERSION = 'v2.5.0-dev'
 SSH_HEADER = 'SSH-{0}-OpenSSH_8.2'  # SSH software to impersonate
 GITHUB_ISSUES_URL = 'https://github.com/jtesta/ssh-audit/issues'  # The URL to the Github issues tracker.
 WINDOWS_MAN_PAGE = ''
diff --git a/src/ssh_audit/hostkeytest.py b/src/ssh_audit/hostkeytest.py
index fe1bba3..87b31fe 100644
--- a/src/ssh_audit/hostkeytest.py
+++ b/src/ssh_audit/hostkeytest.py
@@ -125,8 +125,8 @@ class HostKeyTest:
 
                 # 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.
-                kex_group.send_init(s)
                 try:
+                    kex_group.send_init(s)
                     host_key = kex_group.recv_reply(s, variable_key_len)
                     if host_key is not None:
                         server_kex.set_host_key(host_key_type, host_key)
diff --git a/src/ssh_audit/kexdh.py b/src/ssh_audit/kexdh.py
index 0d10fc9..9c4e08a 100644
--- a/src/ssh_audit/kexdh.py
+++ b/src/ssh_audit/kexdh.py
@@ -327,7 +327,7 @@ class KexGroupExchange(KexDH):
         s.send_packet()
 
         packet_type, payload = s.read_packet(2)
-        if (packet_type != Protocol.MSG_KEXDH_GEX_GROUP) and (packet_type != Protocol.MSG_DEBUG):  # pylint: disable=consider-using-in
+        if packet_type not in [Protocol.MSG_KEXDH_GEX_GROUP, Protocol.MSG_DEBUG]:
             # TODO: replace with a better exception type.
             raise Exception('Expected MSG_KEXDH_GEX_REPLY (%d), but got %d instead.' % (Protocol.MSG_KEXDH_GEX_REPLY, packet_type))