Improved the monkey start function structure a bit, extracted to functions

Prep work for changing system info collection to modular system

monkey/infection_monkey/monkey.py

@@ -32,11 +32,17 @@ from infection_monkey.telemetry.attack.t1106_telem import T1106Telem
 from common.utils.attack_utils import ScanStatus, UsageEnum
 from infection_monkey.exploit.HostExploiter import HostExploiter
 
+MAX_DEPTH_REACHED_MESSAGE = "Reached max depth, shutting down"
+
 __author__ = 'itamar'
 
 LOG = logging.getLogger(__name__)
 
 
+class PlannedShutdown(Exception):
+    pass
+
+
 class InfectionMonkey(object):
     def __init__(self, args):
         self._keep_running = False
@@ -87,22 +93,18 @@ class InfectionMonkey(object):
                 LOG.debug("Default server: %s is already in command servers list" % self._default_server)
 
     def start(self):
-        LOG.info("Monkey is running...")
+        try:
+            LOG.info("Monkey is starting...")
 
+            LOG.debug("Starting the setup phase.")
             # Sets island's IP and port for monkey to communicate to
-        if not self.set_default_server():
-            return
+            self.set_default_server()
             self.set_default_port()
 
             # Create a dir for monkey files if there isn't one
             create_monkey_dir()
 
-        if WindowsUpgrader.should_upgrade():
-            self._upgrading_to_64 = True
-            self._singleton.unlock()
-            LOG.info("32bit monkey running on 64bit Windows. Upgrading.")
-            WindowsUpgrader.upgrade(self._opts)
-            return
+            self.upgrade_to_64_if_needed()
 
             ControlClient.wakeup(parent=self._parent)
             ControlClient.load_control_config()
@@ -110,9 +112,7 @@ class InfectionMonkey(object):
             if is_windows_os():
                 T1106Telem(ScanStatus.USED, UsageEnum.SINGLETON_WINAPI).send()
 
-        if not WormConfiguration.alive:
-            LOG.info("Marked not alive from configuration")
-            return
+            self.shutdown_by_not_alive_config()
 
             if firewall.is_enabled():
                 firewall.add_firewall_rule()
@@ -124,20 +124,12 @@ class InfectionMonkey(object):
             StateTelem(is_done=False).send()
             TunnelTelem().send()
 
-        if WormConfiguration.collect_system_info:
-            LOG.debug("Calling system info collection")
-            system_info_collector = SystemInfoCollector()
-            system_info = system_info_collector.get_info()
-            SystemInfoTelem(system_info).send()
+            LOG.debug("Starting the post-breach phase.")
+            self.collect_system_info_if_configured()
+            PostBreach().execute_all_configured()
 
-        # Executes post breach actions
-        PostBreach().execute()
-
-        if 0 == WormConfiguration.depth:
-            TraceTelem("Reached max depth, shutting down").send()
-            return
-        else:
-            LOG.debug("Running with depth: %d" % WormConfiguration.depth)
+            LOG.debug("Starting the propagation phase.")
+            self.shutdown_by_max_depth_reached()
 
             for iteration_index in range(WormConfiguration.max_iterations):
                 ControlClient.keepalive()
@@ -224,6 +216,35 @@ class InfectionMonkey(object):
             if monkey_tunnel:
                 monkey_tunnel.stop()
                 monkey_tunnel.join()
+        except PlannedShutdown:
+            LOG.info("A planned shutdown of the Monkey occurred. Logging the reason and finishing execution.")
+            LOG.exception("Planned shutdown, reason:")
+
+    def shutdown_by_max_depth_reached(self):
+        if 0 == WormConfiguration.depth:
+            TraceTelem(MAX_DEPTH_REACHED_MESSAGE).send()
+            raise PlannedShutdown(MAX_DEPTH_REACHED_MESSAGE)
+        else:
+            LOG.debug("Running with depth: %d" % WormConfiguration.depth)
+
+    def collect_system_info_if_configured(self):
+        if WormConfiguration.collect_system_info:
+            LOG.debug("Calling system info collection")
+            system_info_collector = SystemInfoCollector()
+            system_info = system_info_collector.get_info()
+            SystemInfoTelem(system_info).send()
+
+    def shutdown_by_not_alive_config(self):
+        if not WormConfiguration.alive:
+            raise PlannedShutdown("Marked 'not alive' from configuration.")
+
+    def upgrade_to_64_if_needed(self):
+        if WindowsUpgrader.should_upgrade():
+            self._upgrading_to_64 = True
+            self._singleton.unlock()
+            LOG.info("32bit monkey running on 64bit Windows. Upgrading.")
+            WindowsUpgrader.upgrade(self._opts)
+            raise PlannedShutdown("Finished upgrading from 32bit to 64bit.")
 
     def cleanup(self):
         LOG.info("Monkey cleanup started")
@@ -346,9 +367,11 @@ class InfectionMonkey(object):
             self._default_server_port = ''
 
     def set_default_server(self):
+        """
+        Sets the default server for the Monkey to communicate back to.
+        :raises PlannedShutdown if couldn't find the server.
+        """
         if not ControlClient.find_server(default_tunnel=self._default_tunnel):
-            LOG.info("Monkey couldn't find server. Going down.")
-            return False
+            raise PlannedShutdown("Monkey couldn't find server with {} default tunnel.".format(self._default_tunnel))
         self._default_server = WormConfiguration.current_server
         LOG.debug("default server set to: %s" % self._default_server)
-        return True

monkey/infection_monkey/post_breach/post_breach_handler.py

@@ -20,7 +20,7 @@ class PostBreach(object):
         self.os_is_linux = not is_windows_os()
         self.pba_list = self.config_to_pba_list()
 
-    def execute(self):
+    def execute_all_configured(self):
         """
         Executes all post breach actions.
         """