Merge pull request #1976 from guardicore/1962-decouple-should_stop-from-config
Move `should_stop` out of the config
This commit is contained in:
commit
a1656a7121
|
@ -72,8 +72,6 @@ class Configuration(object):
|
||||||
###########################
|
###########################
|
||||||
# monkey config
|
# monkey config
|
||||||
###########################
|
###########################
|
||||||
# sets whether or not the monkey is alive. if false will stop scanning and exploiting
|
|
||||||
should_stop = False
|
|
||||||
|
|
||||||
# depth of propagation
|
# depth of propagation
|
||||||
depth = 2
|
depth = 2
|
||||||
|
|
|
@ -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
|
||||||
|
|
|
@ -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()
|
||||||
|
|
|
@ -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": {
|
||||||
|
|
|
@ -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
|
||||||
|
|
|
@ -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)
|
||||||
|
|
||||||
|
|
|
@ -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
|
||||||
|
|
Loading…
Reference in New Issue