forked from p34709852/monkey
Island: Remove agent controls
Agent controls are being replaced by agent signal events
This commit is contained in:
parent
a49ddf7a4a
commit
b11cd9c5f1
|
@ -88,6 +88,7 @@ class MonkeyIslandClient(object):
|
|||
|
||||
@avoid_race_condition
|
||||
def kill_all_monkeys(self):
|
||||
# TODO change this request, because monkey-control resource got removed
|
||||
response = self.requests.post_json(
|
||||
"api/agent-signals/terminate-all-agents", json={"terminate_time": time.time()}
|
||||
)
|
||||
|
|
|
@ -27,7 +27,6 @@ from monkey_island.cc.resources import (
|
|||
TerminateAllAgents,
|
||||
)
|
||||
from monkey_island.cc.resources.AbstractResource import AbstractResource
|
||||
from monkey_island.cc.resources.agent_controls import StopAgentCheck
|
||||
from monkey_island.cc.resources.attack.attack_report import AttackReport
|
||||
from monkey_island.cc.resources.auth import Authenticate, Register, RegistrationStatus, init_jwt
|
||||
from monkey_island.cc.resources.blackbox.log_blackbox_endpoint import LogBlackboxEndpoint
|
||||
|
@ -199,7 +198,6 @@ def init_restful_endpoints(api: FlaskDIWrapper):
|
|||
api.add_resource(PropagationCredentials)
|
||||
api.add_resource(RemoteRun)
|
||||
api.add_resource(Version)
|
||||
api.add_resource(StopAgentCheck)
|
||||
|
||||
# Resources used by black box tests
|
||||
# API Spec: Fix all the following endpoints, see comments in the resource classes
|
||||
|
|
|
@ -1 +0,0 @@
|
|||
from .agent_controls import AgentControls
|
|
@ -1,8 +0,0 @@
|
|||
from mongoengine import Document, FloatField
|
||||
|
||||
|
||||
# TODO rename to Simulation, add other metadata
|
||||
class AgentControls(Document):
|
||||
|
||||
# Timestamp of the last "kill all agents" command
|
||||
last_stop_all = FloatField(default=None)
|
|
@ -22,10 +22,6 @@ from monkey_island.cc.models.monkey_ttl import MonkeyTtl, create_monkey_ttl_docu
|
|||
from monkey_island.cc.server_utils.consts import DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS
|
||||
|
||||
|
||||
class ParentNotFoundError(Exception):
|
||||
"""Raise when trying to get a parent of monkey that doesn't have one"""
|
||||
|
||||
|
||||
class Monkey(Document):
|
||||
"""
|
||||
This class has 2 main section:
|
||||
|
@ -98,18 +94,6 @@ class Monkey(Document):
|
|||
monkey_is_dead = True
|
||||
return monkey_is_dead
|
||||
|
||||
def has_parent(self):
|
||||
for p in self.parent:
|
||||
if p[0] != self.guid:
|
||||
return True
|
||||
return False
|
||||
|
||||
def get_parent(self):
|
||||
if self.has_parent():
|
||||
return Monkey.objects(guid=self.parent[0][0]).first()
|
||||
else:
|
||||
raise ParentNotFoundError(f"No parent was found for agent with GUID {self.guid}")
|
||||
|
||||
def get_os(self):
|
||||
os = "unknown"
|
||||
if self.description.lower().find("linux") != -1:
|
||||
|
|
|
@ -1 +0,0 @@
|
|||
from .stop_agent_check import StopAgentCheck
|
|
@ -1,11 +0,0 @@
|
|||
from monkey_island.cc.resources.AbstractResource import AbstractResource
|
||||
from monkey_island.cc.services.infection_lifecycle import should_agent_die
|
||||
|
||||
|
||||
class StopAgentCheck(AbstractResource):
|
||||
# API Spec: Rename to AgentStopStatus or something, endpoint for this could be
|
||||
# "/api/agents/<GUID>/stop-status"
|
||||
urls = ["/api/monkey-control/needs-to-stop/<int:monkey_guid>"]
|
||||
|
||||
def get(self, monkey_guid: int):
|
||||
return {"stop_agent": should_agent_die(monkey_guid)}
|
|
@ -4,7 +4,6 @@ from flask import jsonify
|
|||
|
||||
from monkey_island.cc.database import mongo
|
||||
from monkey_island.cc.models import Config
|
||||
from monkey_island.cc.models.agent_controls import AgentControls
|
||||
from monkey_island.cc.models.attack.attack_mitigations import AttackMitigations
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
@ -24,7 +23,6 @@ class Database(object):
|
|||
for x in mongo.db.collection_names()
|
||||
if Database._should_drop(x, reset_config)
|
||||
]
|
||||
Database.init_agent_controls()
|
||||
logger.info("DB was reset")
|
||||
return jsonify(status="OK")
|
||||
|
||||
|
@ -44,10 +42,6 @@ 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())
|
||||
|
|
|
@ -1,7 +1,5 @@
|
|||
import logging
|
||||
|
||||
from monkey_island.cc.models import Monkey
|
||||
from monkey_island.cc.models.agent_controls import AgentControls
|
||||
from monkey_island.cc.services.node import NodeService
|
||||
from monkey_island.cc.services.reporting.report import ReportService
|
||||
from monkey_island.cc.services.reporting.report_generation_synchronisation import (
|
||||
|
@ -12,31 +10,6 @@ from monkey_island.cc.services.reporting.report_generation_synchronisation impor
|
|||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def should_agent_die(guid: int) -> bool:
|
||||
monkey = Monkey.objects(guid=str(guid)).first()
|
||||
return _should_agent_stop(monkey) or _is_monkey_killed_manually(monkey)
|
||||
|
||||
|
||||
def _should_agent_stop(monkey: Monkey) -> bool:
|
||||
if monkey.should_stop:
|
||||
# Only stop the agent once, to allow further runs on that machine
|
||||
monkey.should_stop = False
|
||||
monkey.save()
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _is_monkey_killed_manually(monkey: Monkey) -> bool:
|
||||
terminate_timestamp = AgentControls.objects.first().last_stop_all
|
||||
if terminate_timestamp is None:
|
||||
return False
|
||||
if monkey.has_parent():
|
||||
launch_timestamp = monkey.get_parent().launch_time
|
||||
else:
|
||||
launch_timestamp = monkey.launch_time
|
||||
return int(terminate_timestamp) >= int(launch_timestamp)
|
||||
|
||||
|
||||
def get_completed_steps():
|
||||
is_any_exists = NodeService.is_any_monkey_exists()
|
||||
infection_done = NodeService.is_monkey_finished_running()
|
||||
|
|
|
@ -3,7 +3,7 @@ import uuid
|
|||
|
||||
import pytest
|
||||
|
||||
from monkey_island.cc.models.monkey import Monkey, MonkeyNotFoundError, ParentNotFoundError
|
||||
from monkey_island.cc.models.monkey import Monkey, MonkeyNotFoundError
|
||||
from monkey_island.cc.models.monkey_ttl import MonkeyTtl
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
@ -162,35 +162,3 @@ class TestMonkey:
|
|||
|
||||
cache_info_after_query = Monkey.is_monkey.storage.backend.cache_info()
|
||||
assert cache_info_after_query.hits == 2
|
||||
|
||||
@pytest.mark.usefixtures("uses_database")
|
||||
def test_has_parent(self):
|
||||
monkey_1 = Monkey(guid=str(uuid.uuid4()))
|
||||
monkey_2 = Monkey(guid=str(uuid.uuid4()))
|
||||
monkey_1.parent = [[monkey_2.guid]]
|
||||
monkey_1.save()
|
||||
assert monkey_1.has_parent()
|
||||
|
||||
@pytest.mark.usefixtures("uses_database")
|
||||
def test_has_no_parent(self):
|
||||
monkey_1 = Monkey(guid=str(uuid.uuid4()))
|
||||
monkey_1.parent = [[monkey_1.guid]]
|
||||
monkey_1.save()
|
||||
assert not monkey_1.has_parent()
|
||||
|
||||
@pytest.mark.usefixtures("uses_database")
|
||||
def test_get_parent(self):
|
||||
monkey_1 = Monkey(guid=str(uuid.uuid4()))
|
||||
monkey_2 = Monkey(guid=str(uuid.uuid4()))
|
||||
monkey_1.parent = [[monkey_2.guid]]
|
||||
monkey_1.save()
|
||||
monkey_2.save()
|
||||
assert monkey_1.get_parent().guid == monkey_2.guid
|
||||
|
||||
@pytest.mark.usefixtures("uses_database")
|
||||
def test_get_parent_no_parent(self):
|
||||
monkey_1 = Monkey(guid=str(uuid.uuid4()))
|
||||
monkey_1.parent = [[monkey_1.guid]]
|
||||
monkey_1.save()
|
||||
with pytest.raises(ParentNotFoundError):
|
||||
monkey_1.get_parent()
|
||||
|
|
|
@ -1,26 +1,6 @@
|
|||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
from monkey_island.cc.models import Config, Monkey
|
||||
from monkey_island.cc.models.agent_controls import AgentControls
|
||||
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()
|
||||
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.should_stop = True
|
||||
monkey.save()
|
||||
assert not should_agent_die(monkey.guid)
|
||||
|
||||
|
||||
def create_monkey(launch_time):
|
||||
|
@ -32,80 +12,9 @@ def create_monkey(launch_time):
|
|||
return monkey
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("uses_database")
|
||||
def test_should_agent_die_no_kill_event():
|
||||
monkey = create_monkey(launch_time=3)
|
||||
kill_event = AgentControls()
|
||||
kill_event.save()
|
||||
assert not should_agent_die(monkey.guid)
|
||||
|
||||
|
||||
def create_kill_event(event_time):
|
||||
kill_event = AgentControls(last_stop_all=event_time)
|
||||
kill_event.save()
|
||||
return kill_event
|
||||
|
||||
|
||||
def create_parent(child_monkey, launch_time):
|
||||
monkey_parent = Monkey(guid=str(uuid.uuid4()))
|
||||
child_monkey.parent = [[monkey_parent.guid]]
|
||||
monkey_parent.launch_time = launch_time
|
||||
monkey_parent.save()
|
||||
child_monkey.save()
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("uses_database")
|
||||
def test_was_agent_killed_manually(monkeypatch):
|
||||
monkey = create_monkey(launch_time=2)
|
||||
|
||||
create_kill_event(event_time=3)
|
||||
|
||||
assert should_agent_die(monkey.guid)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("uses_database")
|
||||
def test_agent_killed_on_wakeup(monkeypatch):
|
||||
monkey = create_monkey(launch_time=2)
|
||||
|
||||
create_kill_event(event_time=2)
|
||||
|
||||
assert should_agent_die(monkey.guid)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("uses_database")
|
||||
def test_manual_kill_dont_affect_new_monkeys(monkeypatch):
|
||||
monkey = create_monkey(launch_time=3)
|
||||
|
||||
create_kill_event(event_time=2)
|
||||
|
||||
assert not should_agent_die(monkey.guid)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("uses_database")
|
||||
def test_parent_manually_killed(monkeypatch):
|
||||
monkey = create_monkey(launch_time=3)
|
||||
create_parent(child_monkey=monkey, launch_time=1)
|
||||
|
||||
create_kill_event(event_time=2)
|
||||
|
||||
assert should_agent_die(monkey.guid)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("uses_database")
|
||||
def test_parent_manually_killed_on_wakeup(monkeypatch):
|
||||
monkey = create_monkey(launch_time=3)
|
||||
create_parent(child_monkey=monkey, launch_time=2)
|
||||
|
||||
create_kill_event(event_time=2)
|
||||
|
||||
assert should_agent_die(monkey.guid)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("uses_database")
|
||||
def test_manual_kill_dont_affect_new_monkeys_with_parent(monkeypatch):
|
||||
monkey = create_monkey(launch_time=3)
|
||||
create_parent(child_monkey=monkey, launch_time=2)
|
||||
|
||||
create_kill_event(event_time=1)
|
||||
|
||||
assert not should_agent_die(monkey.guid)
|
||||
|
|
Loading…
Reference in New Issue