Agent: Remove InterruptError and use `if` instead

monkey/infection_monkey/exploit/HostExploiter.py

@@ -103,10 +103,6 @@ class HostExploiter:
             self.exploit_result.error_message = "Exploiter has been interrupted"
         return self.interrupt.is_set()
 
-    class InterruptError(Exception):
-        # Raise when exploiter gets interrupted
-        pass
-
     def post_exploit(self):
         self.set_finish_time()
 

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__)
 
@@ -67,10 +68,9 @@ class PowerShellExploiter(HostExploiter):
 
         auth_options = [get_auth_options(creds, use_ssl) for creds in credentials]
 
-        try:
-            self._client = self._authenticate_via_brute_force(credentials, auth_options)
-        except self.InterruptError:
-            self.exploit_result.error_message = "Exploiter has been interrupted"
+        self._client = self._authenticate_via_brute_force(credentials, auth_options)
+
+        if self.is_interrupted():
             return self.exploit_result
 
         if not self._client:
@@ -79,14 +79,9 @@ class PowerShellExploiter(HostExploiter):
             )
             return self.exploit_result
 
-        self.exploit_result.exploitation_success = True
-
         try:
             self._execute_monkey_agent_on_victim()
             self.exploit_result.propagation_success = True
-        except self.InterruptError:
-            self.exploit_result.error_message = "Exploiter has been interrupted"
-            return self.exploit_result
         except Exception as ex:
             logger.error(f"Failed to propagate to the remote host: {ex}")
             self.exploit_result.error_message = str(ex)
@@ -141,9 +136,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):
-            if self.is_interrupted():
-                raise self.InterruptError
+        for (creds, opts) in interruptable_iter(zip(credentials, auth_options), self.interrupt):
             try:
                 client = PowerShellClient(self.host.ip_addr, creds, opts)
                 client.connect()
@@ -152,7 +145,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(
@@ -176,9 +171,6 @@ class PowerShellExploiter(HostExploiter):
     def _execute_monkey_agent_on_victim(self):
         monkey_path_on_victim = self.options["dropper_target_path_win_64"]
 
-        if self.is_interrupted():
-            raise self.InterruptError()
-
         self._copy_monkey_binary_to_victim(monkey_path_on_victim)
         logger.info("Successfully copied the monkey binary to the victim.")