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

@@ -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
+        # TODO Used in propagation phase to set the default server for the victim
-        self._default_tunnel = self._opts.tunnel
-        self._default_server = self._opts.server
         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()
+        self._monkey_inbound_tunnel = ControlClient.create_control_tunnel()
-        if self._monkey_tunnel:
+        if self._monkey_inbound_tunnel:
-            self._monkey_tunnel.start()
+            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):
+        return 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)."
-            )
 
     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:
+            if self._monkey_inbound_tunnel:
-                self._monkey_tunnel.stop()
+                self._monkey_inbound_tunnel.stop()
-                self._monkey_tunnel.join()
+                self._monkey_inbound_tunnel.join()
 
             if firewall.is_enabled():
                 firewall.remove_firewall_rule()