From 11735b4f89cdc7d70ee05981f708e4099532ef8c Mon Sep 17 00:00:00 2001
From: VakarisZ <vakarisz@yahoo.com>
Date: Wed, 8 Dec 2021 09:54:14 +0200
Subject: [PATCH] Island, Agent: small readability and logging improvements
 related to killing the agents

---
 monkey/infection_monkey/monkey.py                       | 5 ++---
 monkey/monkey_island/cc/models/monkey.py                | 2 +-
 monkey/monkey_island/cc/services/database.py            | 8 ++++++--
 monkey/monkey_island/cc/services/infection_lifecycle.py | 4 ----
 4 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py
index a1a11b925..50b145cc0 100644
--- a/monkey/infection_monkey/monkey.py
+++ b/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()
diff --git a/monkey/monkey_island/cc/models/monkey.py b/monkey/monkey_island/cc/models/monkey.py
index 778606151..8e7dccc98 100644
--- a/monkey/monkey_island/cc/models/monkey.py
+++ b/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"
diff --git a/monkey/monkey_island/cc/services/database.py b/monkey/monkey_island/cc/services/database.py
index 14f7296f4..7aeb1bfcf 100644
--- a/monkey/monkey_island/cc/services/database.py
+++ b/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())
diff --git a/monkey/monkey_island/cc/services/infection_lifecycle.py b/monkey/monkey_island/cc/services/infection_lifecycle.py
index 23207e982..25cb76c47 100644
--- a/monkey/monkey_island/cc/services/infection_lifecycle.py
+++ b/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()