Island: Move `should_stop` out of the config to a standalone field in the Monkey document

This commit is contained in:
Shreya Malviya 2022-06-01 16:42:06 +05:30
parent bcbb005ba0
commit d9ce5bae37
6 changed files with 11 additions and 14 deletions

View File

@ -1,4 +1,4 @@
from mongoengine import BooleanField, EmbeddedDocument from mongoengine import EmbeddedDocument
class Config(EmbeddedDocument): class Config(EmbeddedDocument):
@ -11,6 +11,5 @@ class Config(EmbeddedDocument):
See https://mongoengine-odm.readthedocs.io/apireference.html#mongoengine.FieldDoesNotExist See https://mongoengine-odm.readthedocs.io/apireference.html#mongoengine.FieldDoesNotExist
""" """
should_stop = BooleanField()
meta = {"strict": False} meta = {"strict": False}
pass pass

View File

@ -38,6 +38,7 @@ class Monkey(Document):
# SCHEMA # SCHEMA
guid = StringField(required=True) guid = StringField(required=True)
config = EmbeddedDocumentField("Config") config = EmbeddedDocumentField("Config")
should_stop = BooleanField()
dead = BooleanField() dead = BooleanField()
description = StringField() description = StringField()
hostname = StringField() hostname = StringField()

View File

@ -19,11 +19,6 @@ INTERNAL = {
"title": "Monkey", "title": "Monkey",
"type": "object", "type": "object",
"properties": { "properties": {
"should_stop": {
"type": "boolean",
"default": False,
"description": "Was stop command issued for this monkey",
},
"aws_keys": { "aws_keys": {
"type": "object", "type": "object",
"properties": { "properties": {

View File

@ -17,7 +17,7 @@ logger = logging.getLogger(__name__)
def set_stop_all(time: float): def set_stop_all(time: float):
# This will use Agent and Simulation repositories # This will use Agent and Simulation repositories
for monkey in Monkey.objects(): for monkey in Monkey.objects():
monkey.config.should_stop = True monkey.should_stop = True
monkey.save() monkey.save()
agent_controls = AgentControls.objects.first() agent_controls = AgentControls.objects.first()
agent_controls.last_stop_all = time agent_controls.last_stop_all = time
@ -30,9 +30,9 @@ def should_agent_die(guid: int) -> bool:
def _should_agent_stop(monkey: Monkey) -> 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 # Only stop the agent once, to allow further runs on that machine
monkey.config.should_stop = False monkey.should_stop = False
monkey.save() monkey.save()
return True return True
return False return False

View File

@ -261,7 +261,7 @@ class NodeService:
# Cancel the force kill once monkey died # Cancel the force kill once monkey died
if is_dead: 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) mongo.db.monkey.update({"guid": monkey["guid"]}, {"$set": props_to_set}, upsert=False)

View File

@ -10,21 +10,23 @@ from monkey_island.cc.services.infection_lifecycle import should_agent_die
@pytest.mark.usefixtures("uses_database") @pytest.mark.usefixtures("uses_database")
def test_should_agent_die_by_config(monkeypatch): def test_should_agent_die_by_config(monkeypatch):
monkey = Monkey(guid=str(uuid.uuid4())) monkey = Monkey(guid=str(uuid.uuid4()))
monkey.config = Config(should_stop=True) monkey.config = Config()
monkey.should_stop = True
monkey.save() monkey.save()
assert should_agent_die(monkey.guid) assert should_agent_die(monkey.guid)
monkeypatch.setattr( monkeypatch.setattr(
"monkey_island.cc.services.infection_lifecycle._is_monkey_killed_manually", lambda _: False "monkey_island.cc.services.infection_lifecycle._is_monkey_killed_manually", lambda _: False
) )
monkey.config.should_stop = True monkey.should_stop = True
monkey.save() monkey.save()
assert not should_agent_die(monkey.guid) assert not should_agent_die(monkey.guid)
def create_monkey(launch_time): def create_monkey(launch_time):
monkey = Monkey(guid=str(uuid.uuid4())) 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.launch_time = launch_time
monkey.save() monkey.save()
return monkey return monkey