Agent: Refactor the new start and cleanup function

monkey/infection_monkey/monkey.py

@@ -55,7 +55,7 @@ logger = logging.getLogger(__name__)
 class InfectionMonkey(object):
     def __init__(self, args):
         logger.info("Monkey is initializing...")
-        self.master = MockMaster(MockPuppet(), LegacyTelemetryMessengerAdapter())
+        self._master = MockMaster(MockPuppet(), LegacyTelemetryMessengerAdapter())
         self._keep_running = False
         self._exploited_machines = set()
         self._fail_exploitation_machines = set()
@@ -107,12 +107,18 @@ class InfectionMonkey(object):
                 )
 
     def start(self):
-        if not self._is_another_monkey_running():
+        if self._is_another_monkey_running():
+            logger.info("Another instance of the monkey is already running")
+            return
 
         logger.info("Monkey is starting...")
 
         self._connect_to_island()
 
+        # TODO: Reevaluate who is responsible to send this information
+        if is_windows_os():
+            T1106Telem(ScanStatus.USED, UsageEnum.SINGLETON_WINAPI).send()
+
         if InfectionMonkey._is_monkey_alive_by_config():
             logger.error("Monkey marked 'not alive' from configuration.")
             return
@@ -122,9 +128,7 @@ class InfectionMonkey(object):
             return
 
         self._setup()
-            self.master.start()
+        self._master.start()
-        else:
-            logger.info("Another instance of the monkey is already running")
 
     def legacy_start(self):
         if self._is_another_monkey_running():
@@ -193,8 +197,10 @@ class InfectionMonkey(object):
     def _setup(self):
         logger.debug("Starting the setup phase.")
 
-        # Create a dir for monkey files if there isn't one
+        if is_running_on_island():
-        create_monkey_dir()
+            # TODO: Evaluate also this with ControlClient.should_monkey_run
+            # WormConfiguration.started_on_island = True
+            ControlClient.report_start_on_island()
 
         # TODO: Evaluate should we run this check
         # if not ControlClient.should_monkey_run(self._opts.vulnerable_port):
@@ -202,14 +208,6 @@ class InfectionMonkey(object):
         #                  "(it will be exploited later with more depth).")
         #     return False
 
-        # Singleton should handle sending this information
-        if is_windows_os():
-            T1106Telem(ScanStatus.USED, UsageEnum.SINGLETON_WINAPI).send()
-
-        if is_running_on_island():
-            WormConfiguration.started_on_island = True
-            ControlClient.report_start_on_island()
-
         if firewall.is_enabled():
             firewall.add_firewall_rule()
 
@@ -220,9 +218,7 @@ class InfectionMonkey(object):
         StateTelem(is_done=False, version=get_version()).send()
         TunnelTelem().send()
 
-        register_signal_handlers(self.master)
+        register_signal_handlers(self._master)
-
-        return True
 
     def _legacy_setup(self):
         logger.debug("Starting the setup phase.")
@@ -239,23 +235,23 @@ class InfectionMonkey(object):
         ControlClient.wakeup(parent=self._parent)
         ControlClient.load_control_config()
 
+        if is_windows_os():
+            T1106Telem(ScanStatus.USED, UsageEnum.SINGLETON_WINAPI).send()
+
         InfectionMonkey._legacy_shutdown_by_not_alive_config()
 
         self._legacy_upgrade_to_64_if_needed()
 
+        if is_running_on_island():
+            WormConfiguration.started_on_island = True
+            ControlClient.report_start_on_island()
+
         if not ControlClient.should_monkey_run(self._opts.vulnerable_port):
             raise PlannedShutdownException(
                 "Monkey shouldn't run on current machine "
                 "(it will be exploited later with more depth)."
             )
 
-        if is_windows_os():
-            T1106Telem(ScanStatus.USED, UsageEnum.SINGLETON_WINAPI).send()
-
-        if is_running_on_island():
-            WormConfiguration.started_on_island = True
-            ControlClient.report_start_on_island()
-
         if firewall.is_enabled():
             firewall.add_firewall_rule()
 
@@ -522,42 +518,43 @@ class InfectionMonkey(object):
             InfectionMonkey._send_log()
             self._singleton.unlock()
 
-        # self.master.terminate()
-        # self.master.cleanup()
-
         InfectionMonkey._self_delete()
         logger.info("Monkey is shutting down")
 
     def cleanup(self):
         logger.info("Monkey cleanup started")
-        self._keep_running = False
+        try:
+            if self._is_upgrade_to_64_needed():
+                logger.debug("Detected upgrade to 64bit")
+                return
+
+            if self._master:
+                self._master.cleanup()
+
             if self._monkey_tunnel:
                 self._monkey_tunnel.stop()
                 self._monkey_tunnel.join()
 
-        if self._post_breach_phase:
-            self._post_breach_phase.join()
-
             if firewall.is_enabled():
                 firewall.remove_firewall_rule()
                 firewall.close()
 
-        if self._upgrading_to_64:
+            InfectionMonkey._self_delete()
-            InfectionMonkey._close_tunnel()
+
-            firewall.close()
+            InfectionMonkey._send_log()
-        else:
+
             StateTelem(
                 is_done=True, version=get_version()
             ).send()  # Signal the server (before closing the tunnel)
+
+            # TODO: Determine how long between when we
+            #  send telemetry and the monkey actually exits
             InfectionMonkey._close_tunnel()
-            firewall.close()
-            InfectionMonkey._send_log()
             self._singleton.unlock()
-
+        except Exception as e:
-        self.master.terminate()
+            logger.error(f"An error occurred while cleaning up the monkey agent: {e}")
-        self.master.cleanup()
-
             InfectionMonkey._self_delete()
+
         logger.info("Monkey is shutting down")
 
     @staticmethod