From 0ffe023a9f19d4df17767a714e8609012864cbca Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Thu, 17 Mar 2022 14:26:07 -0400
Subject: [PATCH] Agent: Add a query timeout to pymssql.connect()

---
 monkey/infection_monkey/exploit/mssqlexec.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/monkey/infection_monkey/exploit/mssqlexec.py b/monkey/infection_monkey/exploit/mssqlexec.py
index bdef41784..01cc8b59b 100644
--- a/monkey/infection_monkey/exploit/mssqlexec.py
+++ b/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"