From b11cd9c5f17843cc8c94e40dc4f4386f78a9e11b Mon Sep 17 00:00:00 2001 From: vakarisz Date: Fri, 23 Sep 2022 12:08:24 +0300 Subject: [PATCH] Island: Remove agent controls Agent controls are being replaced by agent signal events --- .../island_client/monkey_island_client.py | 1 + monkey/monkey_island/cc/app.py | 2 - .../cc/models/agent_controls/__init__.py | 1 - .../models/agent_controls/agent_controls.py | 8 -- monkey/monkey_island/cc/models/monkey.py | 16 ---- .../cc/resources/agent_controls/__init__.py | 1 - .../agent_controls/stop_agent_check.py | 11 --- monkey/monkey_island/cc/services/database.py | 6 -- .../cc/services/infection_lifecycle.py | 27 ------ .../monkey_island/cc/models/test_monkey.py | 34 +------ .../cc/services/test_infection_lifecycle.py | 91 ------------------- 11 files changed, 2 insertions(+), 196 deletions(-) delete mode 100644 monkey/monkey_island/cc/models/agent_controls/__init__.py delete mode 100644 monkey/monkey_island/cc/models/agent_controls/agent_controls.py delete mode 100644 monkey/monkey_island/cc/resources/agent_controls/__init__.py delete mode 100644 monkey/monkey_island/cc/resources/agent_controls/stop_agent_check.py diff --git a/envs/monkey_zoo/blackbox/island_client/monkey_island_client.py b/envs/monkey_zoo/blackbox/island_client/monkey_island_client.py index f1df4a25b..001a8d17a 100644 --- a/envs/monkey_zoo/blackbox/island_client/monkey_island_client.py +++ b/envs/monkey_zoo/blackbox/island_client/monkey_island_client.py @@ -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()} ) diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index 86616b596..494f48dd0 100644 --- a/monkey/monkey_island/cc/app.py +++ b/monkey/monkey_island/cc/app.py @@ -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 diff --git a/monkey/monkey_island/cc/models/agent_controls/__init__.py b/monkey/monkey_island/cc/models/agent_controls/__init__.py deleted file mode 100644 index e623955c3..000000000 --- a/monkey/monkey_island/cc/models/agent_controls/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .agent_controls import AgentControls diff --git a/monkey/monkey_island/cc/models/agent_controls/agent_controls.py b/monkey/monkey_island/cc/models/agent_controls/agent_controls.py deleted file mode 100644 index 93caa8433..000000000 --- a/monkey/monkey_island/cc/models/agent_controls/agent_controls.py +++ /dev/null @@ -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) diff --git a/monkey/monkey_island/cc/models/monkey.py b/monkey/monkey_island/cc/models/monkey.py index f1732ed66..2fc50681b 100644 --- a/monkey/monkey_island/cc/models/monkey.py +++ b/monkey/monkey_island/cc/models/monkey.py @@ -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: diff --git a/monkey/monkey_island/cc/resources/agent_controls/__init__.py b/monkey/monkey_island/cc/resources/agent_controls/__init__.py deleted file mode 100644 index 4bc6d5b48..000000000 --- a/monkey/monkey_island/cc/resources/agent_controls/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .stop_agent_check import StopAgentCheck diff --git a/monkey/monkey_island/cc/resources/agent_controls/stop_agent_check.py b/monkey/monkey_island/cc/resources/agent_controls/stop_agent_check.py deleted file mode 100644 index 18195e1ce..000000000 --- a/monkey/monkey_island/cc/resources/agent_controls/stop_agent_check.py +++ /dev/null @@ -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//stop-status" - urls = ["/api/monkey-control/needs-to-stop/"] - - def get(self, monkey_guid: int): - return {"stop_agent": should_agent_die(monkey_guid)} diff --git a/monkey/monkey_island/cc/services/database.py b/monkey/monkey_island/cc/services/database.py index fd99f8045..99622aa12 100644 --- a/monkey/monkey_island/cc/services/database.py +++ b/monkey/monkey_island/cc/services/database.py @@ -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()) diff --git a/monkey/monkey_island/cc/services/infection_lifecycle.py b/monkey/monkey_island/cc/services/infection_lifecycle.py index ef11ce5b1..9ada60608 100644 --- a/monkey/monkey_island/cc/services/infection_lifecycle.py +++ b/monkey/monkey_island/cc/services/infection_lifecycle.py @@ -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() diff --git a/monkey/tests/unit_tests/monkey_island/cc/models/test_monkey.py b/monkey/tests/unit_tests/monkey_island/cc/models/test_monkey.py index e25871378..f5a00e5e7 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/models/test_monkey.py +++ b/monkey/tests/unit_tests/monkey_island/cc/models/test_monkey.py @@ -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() 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 10a7f9ade..91f70fb4b 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 @@ -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)