From 3f2fdbaa3ddc3607eaf46b11e2be588ce415f6a1 Mon Sep 17 00:00:00 2001
From: Joe Testa <jtesta@positronsecurity.com>
Date: Tue, 11 Jul 2023 11:08:42 -0400
Subject: [PATCH] Fixed crash during GEX tests.

---
 README.md              |  1 +
 src/ssh_audit/kexdh.py | 24 ++++++++++++++----------
 2 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index e71d319..ca44d2d 100644
--- a/README.md
+++ b/README.md
@@ -182,6 +182,7 @@ For convenience, a web front-end on top of the command-line tool is available at
  - Results from concurrent scans against multiple hosts are no longer improperly combined; bug discovered by [Adam Russell](https://github.com/thecliguy).
  - Hostname resolution failure no longer causes scans against multiple hosts to terminate unexpectedly; credit [Dani Cuesta](https://github.com/daniel-cues).
  - Algorithm recommendations resulting from warnings are now printed in yellow instead of red; credit [Adam Russell](https://github.com/thecliguy).
+ - Fixed crash during GEX tests.
  - Added 1 new key exchange: `curve448-sha512@libssh.org`.
 
 ### v2.9.0 (2023-04-29)
diff --git a/src/ssh_audit/kexdh.py b/src/ssh_audit/kexdh.py
index 9baab73..cdefd4e 100644
--- a/src/ssh_audit/kexdh.py
+++ b/src/ssh_audit/kexdh.py
@@ -26,6 +26,7 @@ import binascii
 import os
 import random
 import struct
+import traceback
 
 # pylint: disable=unused-import
 from typing import Dict, List, Set, Sequence, Tuple, Iterable  # noqa: F401
@@ -375,19 +376,22 @@ class KexGroupExchange(KexDH):
         while packet_type == Protocol.MSG_DEBUG:
             packet_type, payload = s.read_packet(2)
 
-        # Parse the modulus (p) and generator (g) values from the server.
-        ptr = 0
-        p_len = struct.unpack('>I', payload[ptr:ptr + 4])[0]
-        ptr += 4
+        try:
+            # Parse the modulus (p) and generator (g) values from the server.
+            ptr = 0
+            p_len = struct.unpack('>I', payload[ptr:ptr + 4])[0]
+            ptr += 4
 
-        p = int(binascii.hexlify(payload[ptr:ptr + p_len]), 16)
-        ptr += p_len
+            p = int(binascii.hexlify(payload[ptr:ptr + p_len]), 16)
+            ptr += p_len
 
-        g_len = struct.unpack('>I', payload[ptr:ptr + 4])[0]
-        ptr += 4
+            g_len = struct.unpack('>I', payload[ptr:ptr + 4])[0]
+            ptr += 4
 
-        g = int(binascii.hexlify(payload[ptr:ptr + g_len]), 16)
-        ptr += g_len
+            g = int(binascii.hexlify(payload[ptr:ptr + g_len]), 16)
+            ptr += g_len
+        except struct.error:
+            raise KexDHException("Error while parsing modulus and generator during GEX init: %s" % str(traceback.format_exc())) from None
 
         # Now that we got the generator and modulus, perform the DH exchange
         # like usual.