Agent: Make powershell exploiter interruptable

monkey/infection_monkey/exploit/HostExploiter.py

@@ -103,6 +103,10 @@ 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

@@ -67,7 +67,11 @@ 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)
+        try:
+            self._client = self._authenticate_via_brute_force(credentials, auth_options)
+        except self.InterruptError:
+            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"
@@ -79,6 +83,8 @@ class PowerShellExploiter(HostExploiter):
         try:
             self._execute_monkey_agent_on_victim()
             self.exploit_result.propagation_success = True
+        except self.InterruptError:
+            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)
@@ -134,6 +140,8 @@ class PowerShellExploiter(HostExploiter):
         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
             try:
                 client = PowerShellClient(self.host.ip_addr, creds, opts)
                 client.connect()
@@ -166,6 +174,9 @@ 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.")