From d9ce5bae3740324bcbcba69224aadb314561b1dc Mon Sep 17 00:00:00 2001
From: Shreya Malviya <shreya.malviya@gmail.com>
Date: Wed, 1 Jun 2022 16:42:06 +0530
Subject: [PATCH] Island: Move `should_stop` out of the config to a standalone
 field in the Monkey document

---
 monkey/monkey_island/cc/models/config.py                  | 3 +--
 monkey/monkey_island/cc/models/monkey.py                  | 1 +
 .../monkey_island/cc/services/config_schema/internal.py   | 5 -----
 monkey/monkey_island/cc/services/infection_lifecycle.py   | 6 +++---
 monkey/monkey_island/cc/services/node.py                  | 2 +-
 .../monkey_island/cc/services/test_infection_lifecycle.py | 8 +++++---
 6 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/monkey/monkey_island/cc/models/config.py b/monkey/monkey_island/cc/models/config.py
index 0e25b9ee3..febc3e239 100644
--- a/monkey/monkey_island/cc/models/config.py
+++ b/monkey/monkey_island/cc/models/config.py
@@ -1,4 +1,4 @@
-from mongoengine import BooleanField, EmbeddedDocument
+from mongoengine import EmbeddedDocument
 
 
 class Config(EmbeddedDocument):
@@ -11,6 +11,5 @@ class Config(EmbeddedDocument):
     See https://mongoengine-odm.readthedocs.io/apireference.html#mongoengine.FieldDoesNotExist
     """
 
-    should_stop = BooleanField()
     meta = {"strict": False}
     pass
diff --git a/monkey/monkey_island/cc/models/monkey.py b/monkey/monkey_island/cc/models/monkey.py
index bafb9992b..8f1db8121 100644
--- a/monkey/monkey_island/cc/models/monkey.py
+++ b/monkey/monkey_island/cc/models/monkey.py
@@ -38,6 +38,7 @@ class Monkey(Document):
     # SCHEMA
     guid = StringField(required=True)
     config = EmbeddedDocumentField("Config")
+    should_stop = BooleanField()
     dead = BooleanField()
     description = StringField()
     hostname = StringField()
diff --git a/monkey/monkey_island/cc/services/config_schema/internal.py b/monkey/monkey_island/cc/services/config_schema/internal.py
index 86a0089ec..b76db872b 100644
--- a/monkey/monkey_island/cc/services/config_schema/internal.py
+++ b/monkey/monkey_island/cc/services/config_schema/internal.py
@@ -19,11 +19,6 @@ INTERNAL = {
             "title": "Monkey",
             "type": "object",
             "properties": {
-                "should_stop": {
-                    "type": "boolean",
-                    "default": False,
-                    "description": "Was stop command issued for this monkey",
-                },
                 "aws_keys": {
                     "type": "object",
                     "properties": {
diff --git a/monkey/monkey_island/cc/services/infection_lifecycle.py b/monkey/monkey_island/cc/services/infection_lifecycle.py
index 2da5e90b7..84faa1953 100644
--- a/monkey/monkey_island/cc/services/infection_lifecycle.py
+++ b/monkey/monkey_island/cc/services/infection_lifecycle.py
@@ -17,7 +17,7 @@ logger = logging.getLogger(__name__)
 def set_stop_all(time: float):
     # This will use Agent and Simulation repositories
     for monkey in Monkey.objects():
-        monkey.config.should_stop = True
+        monkey.should_stop = True
         monkey.save()
     agent_controls = AgentControls.objects.first()
     agent_controls.last_stop_all = time
@@ -30,9 +30,9 @@ def should_agent_die(guid: int) -> bool:
 
 
 def _should_agent_stop(monkey: Monkey) -> bool:
-    if monkey.config.should_stop:
+    if monkey.should_stop:
         # Only stop the agent once, to allow further runs on that machine
-        monkey.config.should_stop = False
+        monkey.should_stop = False
         monkey.save()
         return True
     return False
diff --git a/monkey/monkey_island/cc/services/node.py b/monkey/monkey_island/cc/services/node.py
index b53696433..4ffb3935d 100644
--- a/monkey/monkey_island/cc/services/node.py
+++ b/monkey/monkey_island/cc/services/node.py
@@ -261,7 +261,7 @@ class NodeService:
 
         # Cancel the force kill once monkey died
         if is_dead:
-            props_to_set["config.should_stop"] = False
+            props_to_set["should_stop"] = False
 
         mongo.db.monkey.update({"guid": monkey["guid"]}, {"$set": props_to_set}, upsert=False)
 
diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_infection_lifecycle.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_infection_lifecycle.py
index 541124248..10a7f9ade 100644
--- a/monkey/tests/unit_tests/monkey_island/cc/services/test_infection_lifecycle.py
+++ b/monkey/tests/unit_tests/monkey_island/cc/services/test_infection_lifecycle.py
@@ -10,21 +10,23 @@ from monkey_island.cc.services.infection_lifecycle import should_agent_die
 @pytest.mark.usefixtures("uses_database")
 def test_should_agent_die_by_config(monkeypatch):
     monkey = Monkey(guid=str(uuid.uuid4()))
-    monkey.config = Config(should_stop=True)
+    monkey.config = Config()
+    monkey.should_stop = True
     monkey.save()
     assert should_agent_die(monkey.guid)
 
     monkeypatch.setattr(
         "monkey_island.cc.services.infection_lifecycle._is_monkey_killed_manually", lambda _: False
     )
-    monkey.config.should_stop = True
+    monkey.should_stop = True
     monkey.save()
     assert not should_agent_die(monkey.guid)
 
 
 def create_monkey(launch_time):
     monkey = Monkey(guid=str(uuid.uuid4()))
-    monkey.config = Config(should_stop=False)
+    monkey.config = Config()
+    monkey.should_stop = False
     monkey.launch_time = launch_time
     monkey.save()
     return monkey