Merge pull request #1790 from guardicore/1611-interruptible-mssql

1611 interruptible mssql

monkey/infection_monkey/exploit/mssqlexec.py

@@ -15,6 +15,7 @@ from infection_monkey.model import DROPPER_ARG
 from infection_monkey.transport import LockedHTTPServer
 from infection_monkey.utils.brute_force import generate_identity_secret_pairs
 from infection_monkey.utils.commands import build_monkey_commandline
+from infection_monkey.utils.threading import interruptable_iter
 
 logger = logging.getLogger(__name__)
 
@@ -22,7 +23,8 @@ logger = logging.getLogger(__name__)
 class MSSQLExploiter(HostExploiter):
     _EXPLOITED_SERVICE = "MSSQL"
     _TARGET_OS_TYPE = ["windows"]
-    LOGIN_TIMEOUT = 15
+    LOGIN_TIMEOUT = LONG_REQUEST_TIMEOUT
+    QUERY_TIMEOUT = LONG_REQUEST_TIMEOUT
     # Time in seconds to wait between MSSQL queries.
     QUERY_BUFFER = 0.5
     SQL_DEFAULT_TCP_PORT = "1433"
@@ -71,6 +73,9 @@ class MSSQLExploiter(HostExploiter):
             )
             return self.exploit_result
 
+        if self.is_interrupted():
+            return self.exploit_result
+
         try:
             # Create dir for payload
             self.create_temp_dir()
@@ -208,12 +213,24 @@ class MSSQLExploiter(HostExploiter):
         """
         # Main loop
         # Iterates on users list
-        for user, password in users_passwords_pairs_list:
+        credentials_iterator = interruptable_iter(
+            users_passwords_pairs_list,
+            self.interrupt,
+            "MSSQL exploiter has been interrupted",
+            logging.INFO,
+        )
+
+        for user, password in credentials_iterator:
             try:
                 # 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"