Agent: Extract propagation loop into a function

monkey/infection_monkey/monkey.py

@@ -62,6 +62,7 @@ class InfectionMonkey(object):
         self._default_server_port = None
         self._opts = None
         self._upgrading_to_64 = False
+        self._monkey_tunnel = None
 
     def initialize(self):
         LOG.info("Monkey is initializing...")
@@ -135,9 +136,9 @@ class InfectionMonkey(object):
             if firewall.is_enabled():
                 firewall.add_firewall_rule()
 
-            monkey_tunnel = ControlClient.create_control_tunnel()
+            self._monkey_tunnel = ControlClient.create_control_tunnel()
-            if monkey_tunnel:
+            if self._monkey_tunnel:
-                monkey_tunnel.start()
+                self._monkey_tunnel.start()
 
             StateTelem(is_done=False, version=get_version()).send()
             TunnelTelem().send()
@@ -149,6 +150,58 @@ class InfectionMonkey(object):
             LOG.debug("Starting the propagation phase.")
             self.shutdown_by_max_depth_reached()
 
+            self.propagate()
+
+            InfectionMonkey.run_ransomware()
+
+            # if host was exploited, before continue to closing the tunnel ensure the exploited
+            # host had its chance to
+            # connect to the tunnel
+            if len(self._exploited_machines) > 0:
+                time_to_sleep = WormConfiguration.keep_tunnel_open_time
+                LOG.info(
+                    "Sleeping %d seconds for exploited machines to connect to tunnel", time_to_sleep
+                )
+                time.sleep(time_to_sleep)
+
+            if self._monkey_tunnel:
+                self._monkey_tunnel.stop()
+                self._monkey_tunnel.join()
+
+            post_breach_phase.join()
+
+        except PlannedShutdownException:
+            LOG.info(
+                "A planned shutdown of the Monkey occurred. Logging the reason and finishing "
+                "execution."
+            )
+            LOG.exception("Planned shutdown, reason:")
+
+    def start_post_breach_phase(self):
+        self.collect_system_info_if_configured()
+        PostBreach().execute_all_configured()
+
+    def shutdown_by_max_depth_reached(self):
+        if 0 == WormConfiguration.depth:
+            TraceTelem(MAX_DEPTH_REACHED_MESSAGE).send()
+            raise PlannedShutdownException(MAX_DEPTH_REACHED_MESSAGE)
+        else:
+            LOG.debug("Running with depth: %d" % WormConfiguration.depth)
+
+    def collect_system_info_if_configured(self):
+        LOG.debug("Calling for system info collection")
+        try:
+            system_info_collector = SystemInfoCollector()
+            system_info = system_info_collector.get_info()
+            SystemInfoTelem(system_info).send()
+        except Exception as e:
+            LOG.exception(f"Exception encountered during system info collection: {str(e)}")
+
+    def shutdown_by_not_alive_config(self):
+        if not WormConfiguration.alive:
+            raise PlannedShutdownException("Marked 'not alive' from configuration.")
+
+    def propagate(self):
         for iteration_index in range(WormConfiguration.max_iterations):
             ControlClient.keepalive()
             ControlClient.load_control_config()
@@ -200,17 +253,13 @@ class InfectionMonkey(object):
                         LOG.debug("Skipping %r - exploitation failed before", machine)
                         continue
 
-                    if monkey_tunnel:
+                if self._monkey_tunnel:
-                        monkey_tunnel.set_tunnel_for_host(machine)
+                    self._monkey_tunnel.set_tunnel_for_host(machine)
                 if self._default_server:
                     if self._network.on_island(self._default_server):
                         machine.set_default_server(
                             get_interface_to_target(machine.ip_addr)
-                                + (
+                            + (":" + self._default_server_port if self._default_server_port else "")
-                                    ":" + self._default_server_port
-                                    if self._default_server_port
-                                    else ""
-                                )
                         )
                     else:
                         machine.set_default_server(self._default_server)
@@ -242,60 +291,11 @@ class InfectionMonkey(object):
                 LOG.info("Sleeping %d seconds before next life cycle iteration", time_to_sleep)
                 time.sleep(time_to_sleep)
 
-            InfectionMonkey.run_ransomware()
-
         if self._keep_running and WormConfiguration.alive:
             LOG.info("Reached max iterations (%d)", WormConfiguration.max_iterations)
         elif not WormConfiguration.alive:
             LOG.info("Marked not alive from configuration")
 
-            # if host was exploited, before continue to closing the tunnel ensure the exploited
-            # host had its chance to
-            # connect to the tunnel
-            if len(self._exploited_machines) > 0:
-                time_to_sleep = WormConfiguration.keep_tunnel_open_time
-                LOG.info(
-                    "Sleeping %d seconds for exploited machines to connect to tunnel", time_to_sleep
-                )
-                time.sleep(time_to_sleep)
-
-            if monkey_tunnel:
-                monkey_tunnel.stop()
-                monkey_tunnel.join()
-
-            post_breach_phase.join()
-
-        except PlannedShutdownException:
-            LOG.info(
-                "A planned shutdown of the Monkey occurred. Logging the reason and finishing "
-                "execution."
-            )
-            LOG.exception("Planned shutdown, reason:")
-
-    def start_post_breach_phase(self):
-        self.collect_system_info_if_configured()
-        PostBreach().execute_all_configured()
-
-    def shutdown_by_max_depth_reached(self):
-        if 0 == WormConfiguration.depth:
-            TraceTelem(MAX_DEPTH_REACHED_MESSAGE).send()
-            raise PlannedShutdownException(MAX_DEPTH_REACHED_MESSAGE)
-        else:
-            LOG.debug("Running with depth: %d" % WormConfiguration.depth)
-
-    def collect_system_info_if_configured(self):
-        LOG.debug("Calling for system info collection")
-        try:
-            system_info_collector = SystemInfoCollector()
-            system_info = system_info_collector.get_info()
-            SystemInfoTelem(system_info).send()
-        except Exception as e:
-            LOG.exception(f"Exception encountered during system info collection: {str(e)}")
-
-    def shutdown_by_not_alive_config(self):
-        if not WormConfiguration.alive:
-            raise PlannedShutdownException("Marked 'not alive' from configuration.")
-
     def upgrade_to_64_if_needed(self):
         if WindowsUpgrader.should_upgrade():
             self._upgrading_to_64 = True