From bcbb005ba049dfd6af2bcd8952321e2e9baad269 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 31 May 2022 17:39:15 +0530 Subject: [PATCH 1/2] Agent: Remove `should_stop` from configuration --- monkey/infection_monkey/config.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/monkey/infection_monkey/config.py b/monkey/infection_monkey/config.py index ddf224b10..41eb8dc93 100644 --- a/monkey/infection_monkey/config.py +++ b/monkey/infection_monkey/config.py @@ -72,8 +72,6 @@ class Configuration(object): ########################### # monkey config ########################### - # sets whether or not the monkey is alive. if false will stop scanning and exploiting - should_stop = False # depth of propagation depth = 2 From d9ce5bae3740324bcbcba69224aadb314561b1dc Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 1 Jun 2022 16:42:06 +0530 Subject: [PATCH 2/2] 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