From ad6b3095237fca8427638a24b23958102fa4ad42 Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Wed, 1 Dec 2021 18:13:27 +0200 Subject: [PATCH] Agent: readability and style changes in monkey.py: refactored back from raising exceptions to logging and returning, not storing part of island config options as separate parameters, etc. --- monkey/infection_monkey/monkey.py | 57 +++++++++++++++---------------- 1 file changed, 28 insertions(+), 29 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index e6733b4a5..a3a0a4f28 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -23,7 +23,6 @@ from infection_monkey.telemetry.messengers.legacy_telemetry_messenger_adapter im from infection_monkey.telemetry.state_telem import StateTelem from infection_monkey.telemetry.tunnel_telem import TunnelTelem from infection_monkey.utils.environment import is_windows_os -from infection_monkey.utils.exceptions.planned_shutdown_error import PlannedShutdownError from infection_monkey.utils.monkey_dir import get_monkey_dir_path, remove_monkey_dir from infection_monkey.utils.monkey_log_path import get_monkey_log_path from infection_monkey.utils.signal_handler import register_signal_handlers @@ -32,24 +31,18 @@ from infection_monkey.windows_upgrader import WindowsUpgrader logger = logging.getLogger(__name__) -class PlannedShutdownError(Exception): - # Raise when we deliberately want to shut down the agent - pass - - class InfectionMonkey: def __init__(self, args): logger.info("Monkey is initializing...") self._master = MockMaster(MockPuppet(), LegacyTelemetryMessengerAdapter()) self._singleton = SystemSingleton() self._opts = self._get_arguments(args) - self._parent = self._opts.parent - self._default_tunnel = self._opts.tunnel - self._default_server = self._opts.server + # TODO Used in propagation phase to set the default server for the victim self._default_server_port = None self._set_propagation_depth() self._add_default_server_to_config() - self._monkey_tunnel = None + # TODO used in propogation phase + self._monkey_inbound_tunnel = None @staticmethod def _get_arguments(args): @@ -87,7 +80,8 @@ class InfectionMonkey: def start(self): if self._is_another_monkey_running(): - raise PlannedShutdownError("Another instance of the monkey is already running.") + logger.info("Another instance of the monkey is already running") + return logger.info("Monkey is starting...") @@ -98,11 +92,13 @@ class InfectionMonkey: T1106Telem(ScanStatus.USED, UsageEnum.SINGLETON_WINAPI).send() if InfectionMonkey._is_monkey_alive_by_config(): - raise PlannedShutdownError("Monkey marked 'not alive' from configuration.") + logger.info("Monkey marked 'not alive' from configuration.") + return if InfectionMonkey._is_upgrade_to_64_needed(): self._upgrade_to_64() - raise PlannedShutdownError("32 bit Agent can't run on 64 bit system.") + logger.info("32 bit Agent can't run on 64 bit system.") + return self._setup() self._master.start() @@ -111,11 +107,13 @@ class InfectionMonkey: # Sets island's IP and port for monkey to communicate to if not self._is_default_server_set(): raise Exception( - "Monkey couldn't find server with {} default tunnel.".format(self._default_tunnel) + "Monkey couldn't find server with {} default tunnel.".format( + self._opts._default_tunnel + ) ) self._set_default_port() - ControlClient.wakeup(parent=self._parent) + ControlClient.wakeup(parent=self._opts._parent) ControlClient.load_control_config() def _is_default_server_set(self) -> bool: @@ -123,7 +121,7 @@ class InfectionMonkey: Sets the default server for the Monkey to communicate back to. :return """ - if not ControlClient.find_server(default_tunnel=self._default_tunnel): + if not ControlClient.find_server(default_tunnel=self._opts._default_tunnel): return False self._default_server = WormConfiguration.current_server logger.debug("default server set to: %s" % self._default_server) @@ -146,21 +144,26 @@ class InfectionMonkey: def _setup(self): logger.debug("Starting the setup phase.") - self._should_run_check_for_performance() + if self._should_exit_for_performance(): + logger.info( + "Monkey shouldn't run on current machine to improve perfomance" + "(it will be exploited later with more depth)." + ) + return if firewall.is_enabled(): firewall.add_firewall_rule() - self._monkey_tunnel = ControlClient.create_control_tunnel() - if self._monkey_tunnel: - self._monkey_tunnel.start() + self._monkey_inbound_tunnel = ControlClient.create_control_tunnel() + if self._monkey_inbound_tunnel: + self._monkey_inbound_tunnel.start() StateTelem(is_done=False, version=get_version()).send() TunnelTelem().send() register_signal_handlers(self._master) - def _should_run_check_for_performance(self): + def _should_exit_for_performance(self): """ This method implements propagation performance enhancing algorithm that kicks in if the run was started from the Island. @@ -170,11 +173,7 @@ class InfectionMonkey: WormConfiguration.started_on_island = True ControlClient.report_start_on_island() - if not ControlClient.should_monkey_run(self._opts.vulnerable_port): - raise PlannedShutdownError( - "Monkey shouldn't run on current machine to improve perfomance" - "(it will be exploited later with more depth)." - ) + return not ControlClient.should_monkey_run(self._opts.vulnerable_port) def _is_another_monkey_running(self): return not self._singleton.try_lock() @@ -195,9 +194,9 @@ class InfectionMonkey: if self._master: self._master.cleanup() - if self._monkey_tunnel: - self._monkey_tunnel.stop() - self._monkey_tunnel.join() + if self._monkey_inbound_tunnel: + self._monkey_inbound_tunnel.stop() + self._monkey_inbound_tunnel.join() if firewall.is_enabled(): firewall.remove_firewall_rule()