Sockets now time out after 30 seconds during connection rate testing.

src/ssh_audit/dheat.py

@@ -51,7 +51,7 @@ class DHEat:
     MAX_SAFE_RATE = 20.0
 
     # The warning added to DH algorithms in the UI when dh_rate_test determines that no throttling is being done.
-    DHEAT_WARNING = "Potentially insufficient connection throttling detected, resulting in possible vulnerability to the DHEat DoS attack (CVE-2002-20001).  Either connection throttling or removal of Diffie-Hellman key exchanges is necessary to remediate this issue.  Suppress this test/message with --skip-rate-test.  Additional info: {connections:d} connections were created in {time_elapsed:.3f} seconds, or {rate:.1f} conns/sec; server must respond with a rate less than {max_safe_rate:.1f} conns/sec to be considered safe."
+    DHEAT_WARNING = "Potentially insufficient connection throttling detected, resulting in possible vulnerability to the DHEat DoS attack (CVE-2002-20001).  Suppress this test and message with the --skip-rate-test option.  Additional info: {connections:d} connections were created in {time_elapsed:.3f} seconds, or {rate:.1f} conns/sec; server must respond with a rate less than {max_safe_rate:.1f} conns/sec to be considered safe."
 
     # List of the Diffie-Hellman group exchange algorithms this test supports.
     gex_algs = [
@@ -309,14 +309,14 @@ class DHEat:
     def _dh_rate_test(out: 'OutputBuffer', aconf: 'AuditConf', kex: 'SSH2_Kex', max_time: float, max_connections: int, concurrent_sockets: int) -> str:
         '''Attempts to quickly create many sockets to the target server.  This simulates the DHEat attack without causing an actual DoS condition.  If a rate greater than MAX_SAFE_RATE is allowed, then a warning string is returned.'''
 
-        def _close_socket(socket_list: List[socket.socket], s: socket.socket) -> None:
+        def _close_socket(socket_dict: Dict[socket.socket, float], s: socket.socket) -> None:
             try:
                 s.shutdown(socket.SHUT_RDWR)
                 s.close()
             except OSError:
                 pass
 
-            socket_list.remove(s)
+            del socket_dict[s]
 
         if sys.platform == "win32":
             DHEat.YELLOWB = "\033[1;93m"
@@ -357,22 +357,25 @@ class DHEat:
                 out.d("Skipping DHEat.dh_rate_test() since server does not support any DH algorithms: [%s]" % ", ".join(kex.kex_algorithms))
                 return ""
             else:
-                out.d("DHEat.dh_rate_test(): starting test; parameters: %f seconds, %u max connections, %u concurrent sockets." % (max_time, max_connections, concurrent_sockets))
+                out.d("DHEat.dh_rate_test(): starting test; parameters: %f seconds, %u max connections, %u concurrent sockets." % (max_time, max_connections, concurrent_sockets), write_now=True)
 
         num_attempted_connections = 0
         num_opened_connections = 0
-        socket_list: List[socket.socket] = []
+        socket_dict: Dict[socket.socket, float] = {}
         start_timer = time.time()
+        now = start_timer
         last_update = start_timer
         while True:
+            now = time.time()
 
             # During non-interactive tests, limit based on time and number of connections.  Otherwise, we loop indefinitely until the user presses CTRL-C.
-            if (interactive is False) and ((time.time() - start_timer) >= max_time) and (num_opened_connections >= max_connections):
+            if (interactive is False) and ((now - start_timer) >= max_time) or (num_opened_connections >= max_connections):
                 break
 
+            # out.d("interactive: %r; time.time() - start_timer: %f; max_time: %f; num_opened_connections: %u; max_connections: %u" % (interactive, time.time() - start_timer, max_time, num_opened_connections, max_connections), write_now=True)
+
             # Give the user some interactive feedback.
             if interactive:
-                now = time.time()
                 if (now - last_update) >= 1.0:
                     seconds_running = now - start_timer
                     print("%s%s%s Run time: %s%.1f%s; TCP SYNs: %s%u%s; Compl. conns: %s%u%s; TCP SYNs/sec: %s%.1f%s; Compl. conns/sec: %s%.1f%s    \r" % (DHEat.WHITEB, spinner[spinner_index], DHEat.CLEAR, DHEat.WHITEB, seconds_running, DHEat.CLEAR, DHEat.WHITEB, num_attempted_connections, DHEat.CLEAR, DHEat.WHITEB, num_opened_connections, DHEat.CLEAR, DHEat.BLUEB, num_attempted_connections / seconds_running, DHEat.CLEAR, DHEat.BLUEB, num_opened_connections / seconds_running, DHEat.CLEAR), end="")
@@ -388,25 +391,48 @@ class DHEat:
                         if sleep_time > 0.0:
                             time.sleep(sleep_time)
 
-            while (len(socket_list) < concurrent_sockets) and (len(socket_list) + num_opened_connections < max_connections):
+            # Check our sockets to see if they've existed for more than 30 seconds.  If so, close them so new ones can be re-opened in their place.
+            timedout_sockets = []
+            for s, create_time in socket_dict.items():
+                if (now - create_time) > 30:
+                    timedout_sockets.append(s)  # We shouldn't modify the dictionary while iterating over it, so add it to a separate list.
+
+            # Now we can safely close the timed-out sockets.
+            while True:
+                if len(timedout_sockets) == 0:  # Ensure that len() is called in every iteration by putting it here instead of the while clause.
+                    break
+
+                out.d("Closing timed-out socket.", write_now=True)
+                _close_socket(socket_dict, timedout_sockets[0])
+                del timedout_sockets[0]
+
+            # Open new sockets until we've hit the number of concurrent sockets, or if we exceeded the number of maximum connections.
+            while (len(socket_dict) < concurrent_sockets) and (len(socket_dict) + num_opened_connections < max_connections):
                 s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                 s.setblocking(False)
 
-                out.d("Creating socket (%u of %u already exist)..." % (len(socket_list), concurrent_sockets))
+                # out.d("Creating socket (%u of %u already exist)..." % (len(socket_dict), concurrent_sockets), write_now=True)
                 ret = s.connect_ex((aconf.host, aconf.port))
                 num_attempted_connections += 1
                 if ret in [0, 115]:  # Check if connection is successful or EINPROGRESS.
-                    socket_list.append(s)
+                    socket_dict[s] = now
+                else:
+                    out.d("connect_ex() returned: %d" % ret, write_now=True)
 
+            # out.d("Calling select() on %u sockets..." % len(socket_dict), write_now=True)
+            socket_list: List[socket.socket] = [*socket_dict]  # Get a list of sockets from the dictionary.
             rlist, _, elist = select.select(socket_list, [], socket_list, 0.1)
 
             # For each socket that has something for us to read...
             for s in rlist:
+                # out.d("Socket in read list.", write_now=True)
                 buf = b''
                 try:
                     buf = s.recv(8)
-                except (ConnectionResetError, BrokenPipeError):
-                    _close_socket(socket_list, s)
+                    # out.d("Banner: %r" % buf, write_now=True)
+                except (ConnectionRefusedError, ConnectionResetError, BrokenPipeError, TimeoutError):
+                    out.d("Socket error.", write_now=True)
+                    _close_socket(socket_dict, s)
                     continue
 
                 # If we received the SSH header, we'll count this as an opened connection.
@@ -414,7 +440,7 @@ class DHEat:
                     num_opened_connections += 1
                     out.d("Number of opened connections: %u (max: %u)." % (num_opened_connections, max_connections))
 
-                _close_socket(socket_list, s)
+                _close_socket(socket_dict, s)
 
                 # Since we just closed the socket, ensure its not in the exception list.
                 if s in elist:
@@ -422,20 +448,24 @@ class DHEat:
 
             # Close all sockets that are in the exception state.
             for s in elist:
-                _close_socket(socket_list, s)
+                # out.d("Socket in exception list.", write_now=True)
+                _close_socket(socket_dict, s)
 
         # Close any remaining sockets.
-        while len(socket_list) > 0:
-            _close_socket(socket_list, socket_list[0])
+        while True:
+            if len(socket_dict) == 0:  # Ensure that len() is called in every iteration by putting it here instead of the while clause.
+                break
+
+            _close_socket(socket_dict, [*socket_dict][0])  # Close & remove the first socket we find.
 
         time_elapsed = time.time() - start_timer
-        out.d("DHEat.dh_rate_test() results: time elapsed: %f; connections created: %u" % (time_elapsed, num_opened_connections))
+        out.d("DHEat.dh_rate_test() results: time elapsed: %f; connections created: %u" % (time_elapsed, num_opened_connections), write_now=True)
 
         note = ""
         rate = 0.0
         if time_elapsed > 0.0 and num_opened_connections > 0:
             rate = num_opened_connections / time_elapsed
-            out.d("DHEat.dh_rate_test() results: %.1f connections opened per second." % rate)
+            out.d("DHEat.dh_rate_test() results: %.1f connections opened per second." % rate, write_now=True)
 
             # If we were able to open connections at a rate greater than 25 per second, then we need to warn the user.
             if rate > DHEat.MAX_SAFE_RATE: