From 83b18debc0e4ac5c1e831cb90ebf38ba9fe98ef1 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 21 Mar 2022 08:39:49 -0400 Subject: [PATCH] Agent: Remove InterruptError and use `if` instead --- .../infection_monkey/exploit/HostExploiter.py | 4 ---- monkey/infection_monkey/exploit/powershell.py | 22 ++++++------------- 2 files changed, 7 insertions(+), 19 deletions(-) diff --git a/monkey/infection_monkey/exploit/HostExploiter.py b/monkey/infection_monkey/exploit/HostExploiter.py index b1e2c72d3..2e198ac4c 100644 --- a/monkey/infection_monkey/exploit/HostExploiter.py +++ b/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() diff --git a/monkey/infection_monkey/exploit/powershell.py b/monkey/infection_monkey/exploit/powershell.py index 9840ac7a7..b097630d7 100644 --- a/monkey/infection_monkey/exploit/powershell.py +++ b/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.")