Agent: Add a query timeout to pymssql.connect()

monkey/infection_monkey/exploit/mssqlexec.py

@@ -23,6 +23,7 @@ class MSSQLExploiter(HostExploiter):
     _EXPLOITED_SERVICE = "MSSQL"
     _TARGET_OS_TYPE = ["windows"]
     LOGIN_TIMEOUT = 15
+    QUERY_TIMEOUT = LOGIN_TIMEOUT
     # Time in seconds to wait between MSSQL queries.
     QUERY_BUFFER = 0.5
     SQL_DEFAULT_TCP_PORT = "1433"
@@ -213,7 +214,12 @@ class MSSQLExploiter(HostExploiter):
                 # Core steps
                 # Trying to connect
                 conn = pymssql.connect(
-                    host, user, password, port=port, login_timeout=self.LOGIN_TIMEOUT
+                    host,
+                    user,
+                    password,
+                    port=port,
+                    login_timeout=self.LOGIN_TIMEOUT,
+                    timeout=self.QUERY_TIMEOUT,
                 )
                 logger.info(
                     f"Successfully connected to host: {host} using user: {user} and password"