Merge pull request #1791 from guardicore/1611-interruptable-powershell

1611 Make powershell exploiter interruptable

monkey/infection_monkey/exploit/HostExploiter.py

@@ -94,14 +94,15 @@ class HostExploiter:
         )
         self.set_start_time()
 
-    def is_interrupted(self):
+    def _is_interrupted(self):
+        return self.interrupt.is_set()
+
+    def _set_interrupted(self):
         # This method should be refactored to raise an exception to reduce duplication in the
         # "if is_interrupted: return self.exploitation_results"
         # Ideally the user should only do "check_for_interrupt()"
-        if self.interrupt.is_set():
         logger.info("Exploiter has been interrupted")
-            self.exploit_result.error_message = "Exploiter has been interrupted"
+        self.exploit_result.interrupted = True
-        return self.interrupt.is_set()
 
     def post_exploit(self):
         self.set_finish_time()

monkey/infection_monkey/exploit/mssqlexec.py

@@ -73,7 +73,8 @@ class MSSQLExploiter(HostExploiter):
             )
             return self.exploit_result
 
-        if self.is_interrupted():
+        if self._is_interrupted():
+            self._set_interrupted()
             return self.exploit_result
 
         try:

monkey/infection_monkey/exploit/powershell.py

@@ -23,6 +23,7 @@ from infection_monkey.exploit.tools.helpers import get_random_file_suffix
 from infection_monkey.model import DROPPER_ARG, RUN_MONKEY, VictimHost
 from infection_monkey.utils.commands import build_monkey_commandline
 from infection_monkey.utils.environment import is_windows_os
+from infection_monkey.utils.threading import interruptable_iter
 
 logger = logging.getLogger(__name__)
 
@@ -68,14 +69,17 @@ class PowerShellExploiter(HostExploiter):
         auth_options = [get_auth_options(creds, use_ssl) for creds in credentials]
 
         self._client = self._authenticate_via_brute_force(credentials, auth_options)
+
+        if self._is_interrupted():
+            self._set_interrupted()
+            return self.exploit_result
+
         if not self._client:
             self.exploit_result.error_message = (
                 "Unable to authenticate to the remote host using any of the available credentials"
             )
             return self.exploit_result
 
-        self.exploit_result.exploitation_success = True
-
         try:
             self._execute_monkey_agent_on_victim()
             self.exploit_result.propagation_success = True
@@ -133,7 +137,7 @@ class PowerShellExploiter(HostExploiter):
     def _authenticate_via_brute_force(
         self, credentials: List[Credentials], auth_options: List[AuthOptions]
     ) -> Optional[IPowerShellClient]:
-        for (creds, opts) in zip(credentials, auth_options):
+        for (creds, opts) in interruptable_iter(zip(credentials, auth_options), self.interrupt):
             try:
                 client = PowerShellClient(self.host.ip_addr, creds, opts)
                 client.connect()
@@ -142,7 +146,9 @@ class PowerShellExploiter(HostExploiter):
                     f"{creds.username}, Secret Type: {creds.secret_type.name}"
                 )
 
+                self.exploit_result.exploitation_success = True
                 self._report_login_attempt(True, creds)
+
                 return client
             except Exception as ex:
                 logger.debug(

monkey/infection_monkey/exploit/wmiexec.py

@@ -70,7 +70,8 @@ class WmiExploiter(HostExploiter):
 
             downloaded_agent = self.agent_repository.get_agent_binary(self.host.os["type"])
 
-            if self.is_interrupted():
+            if self._is_interrupted():
+                self._set_interrupted()
                 return self.exploit_result
 
             remote_full_path = SmbTools.copy_file(

monkey/infection_monkey/i_puppet/i_puppet.py

@@ -24,6 +24,7 @@ class UnknownPluginError(Exception):
 class ExploiterResultData:
     exploitation_success: bool = False
     propagation_success: bool = False
+    interrupted: bool = False
     os: str = ""
     info: Mapping = None
     attempts: Iterable = None

monkey/infection_monkey/puppet/mock_puppet.py

@@ -190,17 +190,18 @@ class MockPuppet(IPuppet):
         successful_exploiters = {
             DOT_1: {
                 "PowerShellExploiter": ExploiterResultData(
-                    True, True, os_windows, info_powershell, attempts, None
+                    True, True, False, os_windows, info_powershell, attempts, None
                 ),
                 "ZerologonExploiter": ExploiterResultData(
-                    False, False, os_windows, {}, [], "Zerologon failed"
+                    False, False, False, os_windows, {}, [], "Zerologon failed"
                 ),
                 "SSHExploiter": ExploiterResultData(
-                    False, False, os_linux, info_ssh, attempts, "Failed exploiting"
+                    False, False, False, os_linux, info_ssh, attempts, "Failed exploiting"
                 ),
             },
             DOT_3: {
                 "PowerShellExploiter": ExploiterResultData(
+                    False,
                     False,
                     False,
                     os_windows,
@@ -209,9 +210,11 @@ class MockPuppet(IPuppet):
                     "PowerShell Exploiter Failed",
                 ),
                 "SSHExploiter": ExploiterResultData(
-                    False, False, os_linux, info_ssh, attempts, "Failed exploiting"
+                    False, False, False, os_linux, info_ssh, attempts, "Failed exploiting"
+                ),
+                "ZerologonExploiter": ExploiterResultData(
+                    True, False, False, os_windows, {}, [], None
                 ),
-                "ZerologonExploiter": ExploiterResultData(True, False, os_windows, {}, [], None),
             },
         }
 
@@ -219,7 +222,7 @@ class MockPuppet(IPuppet):
             return successful_exploiters[host.ip_addr][name]
         except KeyError:
             return ExploiterResultData(
-                False, False, os_linux, {}, [], f"{name} failed for host {host}"
+                False, False, False, os_linux, {}, [], f"{name} failed for host {host}"
             )
 
     def run_payload(self, name: str, options: Dict, interrupt: threading.Event):

monkey/infection_monkey/telemetry/exploit_telem.py

@@ -1,9 +1,9 @@
 from typing import Dict
 
 from common.common_consts.telem_categories import TelemCategoryEnum
+from infection_monkey.i_puppet.i_puppet import ExploiterResultData
 from infection_monkey.model.host import VictimHost
 from infection_monkey.telemetry.base_telem import BaseTelem
-from infection_monkey.i_puppet.i_puppet import ExploiterResultData
 
 
 class ExploitTelem(BaseTelem):
@@ -25,6 +25,7 @@ class ExploitTelem(BaseTelem):
         self.host = host.__dict__
         self.exploitation_result = result.exploitation_success
         self.propagation_result = result.propagation_success
+        self.interrupted = result.interrupted
         self.info = result.info
         self.attempts = result.attempts
 
@@ -34,6 +35,7 @@ class ExploitTelem(BaseTelem):
         return {
             "exploitation_result": self.exploitation_result,
             "propagation_result": self.propagation_result,
+            "interrupted": self.interrupted,
             "machine": self.host,
             "exploiter": self.name,
             "info": self.info,

monkey/tests/unit_tests/infection_monkey/master/test_propagator.py

@@ -201,38 +201,38 @@ class MockExploiter:
                 results_callback(
                     "PowerShellExploiter",
                     host,
-                    ExploiterResultData(True, True, os_windows, {}, {}, None),
+                    ExploiterResultData(True, True, False, os_windows, {}, {}, None),
                 )
                 results_callback(
                     "SSHExploiter",
                     host,
-                    ExploiterResultData(False, False, os_linux, {}, {}, "SSH FAILED for .1"),
+                    ExploiterResultData(False, False, False, os_linux, {}, {}, "SSH FAILED for .1"),
                 )
             elif host.ip_addr.endswith(".2"):
                 results_callback(
                     "PowerShellExploiter",
                     host,
                     ExploiterResultData(
-                        False, False, os_windows, {}, {}, "POWERSHELL FAILED for .2"
+                        False, False, False, os_windows, {}, {}, "POWERSHELL FAILED for .2"
                     ),
                 )
                 results_callback(
                     "SSHExploiter",
                     host,
-                    ExploiterResultData(False, False, os_linux, {}, {}, "SSH FAILED for .2"),
+                    ExploiterResultData(False, False, False, os_linux, {}, {}, "SSH FAILED for .2"),
                 )
             elif host.ip_addr.endswith(".3"):
                 results_callback(
                     "PowerShellExploiter",
                     host,
                     ExploiterResultData(
-                        False, False, os_windows, {}, {}, "POWERSHELL FAILED for .3"
+                        False, False, False, os_windows, {}, {}, "POWERSHELL FAILED for .3"
                     ),
                 )
                 results_callback(
                     "SSHExploiter",
                     host,
-                    ExploiterResultData(True, True, os_linux, {}, {}, None),
+                    ExploiterResultData(True, True, False, os_linux, {}, {}, None),
                 )
 
 

monkey/tests/unit_tests/infection_monkey/telemetry/test_exploit_telem.py

@@ -40,7 +40,7 @@ def exploit_telem_test_instance():
         EXPLOITER_NAME,
         HOST,
         ExploiterResultData(
-            RESULT, RESULT, OS_LINUX, EXPLOITER_INFO, EXPLOITER_ATTEMPTS, ERROR_MSG
+            RESULT, RESULT, False, OS_LINUX, EXPLOITER_INFO, EXPLOITER_ATTEMPTS, ERROR_MSG
         ),
     )
 
@@ -50,6 +50,7 @@ def test_exploit_telem_send(exploit_telem_test_instance, spy_send_telemetry):
     expected_data = {
         "exploitation_result": RESULT,
         "propagation_result": RESULT,
+        "interrupted": False,
         "machine": HOST_AS_DICT,
         "exploiter": EXPLOITER_NAME,
         "info": EXPLOITER_INFO,