Island, Agent: small readability and logging improvements related to killing the agents

monkey/infection_monkey/monkey.py

@@ -74,11 +74,10 @@ class InfectionMonkey:
         if is_windows_os():
             T1106Telem(ScanStatus.USED, UsageEnum.SINGLETON_WINAPI).send()
 
-        # TODO move this function
         should_stop = ControlChannel(WormConfiguration.current_server, GUID).should_agent_stop()
-        logger.info(f"Should monkey stop: {should_stop}")
         if should_stop:
-            sys.exit(1)
+            logger.info("The Monkey Island has instructed this agent to stop.")
+            return
 
         if InfectionMonkey._is_upgrade_to_64_needed():
             self._upgrade_to_64()

monkey/monkey_island/cc/models/monkey.py

@@ -110,7 +110,7 @@ class Monkey(Document):
         if self.has_parent():
             Monkey.objects(guid=self.parent[0][0]).first()
         else:
-            raise ParentNotFoundError
+            raise ParentNotFoundError(f"No parent was found for agent with GUID {self.guid}")
 
     def get_os(self):
         os = "unknown"

monkey/monkey_island/cc/services/database.py

@@ -3,8 +3,8 @@ import logging
 from flask import jsonify
 
 from monkey_island.cc.database import mongo
+from monkey_island.cc.models.agent_controls import AgentControls
 from monkey_island.cc.models.attack.attack_mitigations import AttackMitigations
-from monkey_island.cc.services.infection_lifecycle import init_agent_controls
 from monkey_island.cc.services.config import ConfigService
 
 logger = logging.getLogger(__name__)
@@ -24,7 +24,7 @@ class Database(object):
             if not x.startswith("system.") and not x == AttackMitigations.COLLECTION_NAME
         ]
         ConfigService.init_config()
-        init_agent_controls()
+        Database.init_agent_controls()
         logger.info("DB was reset")
         return jsonify(status="OK")
 
@@ -33,6 +33,10 @@ class Database(object):
         mongo.db[collection_name].drop()
         logger.info("Dropped collection {}".format(collection_name))
 
+    @staticmethod
+    def init_agent_controls():
+        AgentControls().save()
+
     @staticmethod
     def is_mitigations_missing() -> bool:
         return bool(AttackMitigations.COLLECTION_NAME not in mongo.db.list_collection_names())

monkey/monkey_island/cc/services/infection_lifecycle.py

@@ -41,10 +41,6 @@ def _is_monkey_killed_manually(monkey: Monkey) -> bool:
     return int(kill_timestamp) >= int(launch_timestamp)
 
 
-def init_agent_controls():
-    AgentControls().save()
-
-
 def get_completed_steps():
     is_any_exists = NodeService.is_any_monkey_exists()
     infection_done = NodeService.is_monkey_finished_running()