From be30db885bd4385d0f85e986409e7b2259008c0b Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Thu, 22 Sep 2022 11:59:46 +0200 Subject: [PATCH 01/67] Island: Add TERMINATE_AGENTS to IslandEventTopic --- monkey/monkey_island/cc/event_queue/i_island_event_queue.py | 1 + 1 file changed, 1 insertion(+) diff --git a/monkey/monkey_island/cc/event_queue/i_island_event_queue.py b/monkey/monkey_island/cc/event_queue/i_island_event_queue.py index cf123c3e2..7ecf8c01b 100644 --- a/monkey/monkey_island/cc/event_queue/i_island_event_queue.py +++ b/monkey/monkey_island/cc/event_queue/i_island_event_queue.py @@ -9,6 +9,7 @@ class IslandEventTopic(Enum): CLEAR_SIMULATION_DATA = auto() RESET_AGENT_CONFIGURATION = auto() SET_ISLAND_MODE = auto() + TERMINATE_AGENTS = auto() class IIslandEventQueue(ABC): From 560d941885c9a30b10cae3f12d4b20eb7eeb629f Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Thu, 22 Sep 2022 12:00:20 +0200 Subject: [PATCH 02/67] Island: Add terminate field to Simulation model --- monkey/monkey_island/cc/models/simulation.py | 3 +++ vulture_allowlist.py | 4 +++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/models/simulation.py b/monkey/monkey_island/cc/models/simulation.py index d04bee76b..3f52ae370 100644 --- a/monkey/monkey_island/cc/models/simulation.py +++ b/monkey/monkey_island/cc/models/simulation.py @@ -1,6 +1,8 @@ from __future__ import annotations +from datetime import datetime from enum import Enum +from typing import Optional from common.base_models import InfectionMonkeyBaseModel @@ -13,3 +15,4 @@ class IslandMode(Enum): class Simulation(InfectionMonkeyBaseModel): mode: IslandMode = IslandMode.UNSET + terminate_signal_time: Optional[datetime] = None diff --git a/vulture_allowlist.py b/vulture_allowlist.py index ef9613fca..43292e49e 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -11,7 +11,7 @@ from common.credentials import Credentials, LMHash, NTHash from common.types import SocketAddress from infection_monkey.exploit.log4shell_utils.ldap_server import LDAPServerFactory from monkey_island.cc.event_queue import IslandEventTopic, PyPubSubIslandEventQueue -from monkey_island.cc.models import Report +from monkey_island.cc.models import Report, Simulation from monkey_island.cc.models.networkmap import Arc, NetworkMap from monkey_island.cc.repository import MongoAgentRepository, MongoMachineRepository from monkey_island.cc.repository.attack.IMitigationsRepository import IMitigationsRepository @@ -328,3 +328,5 @@ CC_TUNNEL # TODO: Remove after #2323 SocketAddress + +Simulation.terminate_signal_time From 721cc295591f75b06f091b994fa413facf67b987 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Thu, 22 Sep 2022 12:02:47 +0200 Subject: [PATCH 03/67] Island: Add AgentSignals model --- monkey/monkey_island/cc/models/__init__.py | 1 + monkey/monkey_island/cc/models/agent_signals.py | 8 ++++++++ vulture_allowlist.py | 2 +- 3 files changed, 10 insertions(+), 1 deletion(-) create mode 100644 monkey/monkey_island/cc/models/agent_signals.py diff --git a/monkey/monkey_island/cc/models/__init__.py b/monkey/monkey_island/cc/models/__init__.py index 65e63fe14..ca4078faa 100644 --- a/monkey/monkey_island/cc/models/__init__.py +++ b/monkey/monkey_island/cc/models/__init__.py @@ -15,3 +15,4 @@ from .communication_type import CommunicationType from .node import Node from common.types import AgentID from .agent import Agent +from .agent_signals import AgentSignals diff --git a/monkey/monkey_island/cc/models/agent_signals.py b/monkey/monkey_island/cc/models/agent_signals.py new file mode 100644 index 000000000..37af7b4c1 --- /dev/null +++ b/monkey/monkey_island/cc/models/agent_signals.py @@ -0,0 +1,8 @@ +from datetime import datetime +from typing import Optional + +from common.base_models import InfectionMonkeyBaseModel + + +class AgentSignals(InfectionMonkeyBaseModel): + terminate: Optional[datetime] diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 43292e49e..742450c50 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -11,7 +11,7 @@ from common.credentials import Credentials, LMHash, NTHash from common.types import SocketAddress from infection_monkey.exploit.log4shell_utils.ldap_server import LDAPServerFactory from monkey_island.cc.event_queue import IslandEventTopic, PyPubSubIslandEventQueue -from monkey_island.cc.models import Report, Simulation +from monkey_island.cc.models import AgentSignals, Report, Simulation from monkey_island.cc.models.networkmap import Arc, NetworkMap from monkey_island.cc.repository import MongoAgentRepository, MongoMachineRepository from monkey_island.cc.repository.attack.IMitigationsRepository import IMitigationsRepository From c0afae6dfada588b76c695c49d27ab5b5df8aade Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Thu, 22 Sep 2022 12:19:15 +0200 Subject: [PATCH 04/67] Island: Add stubbed AgentSignalsService --- monkey/monkey_island/cc/services/__init__.py | 1 + .../cc/services/agent_signals_service.py | 26 +++++++++++++++++++ vulture_allowlist.py | 21 ++++++++------- 3 files changed, 38 insertions(+), 10 deletions(-) create mode 100644 monkey/monkey_island/cc/services/agent_signals_service.py diff --git a/monkey/monkey_island/cc/services/__init__.py b/monkey/monkey_island/cc/services/__init__.py index d75734d59..c73aff356 100644 --- a/monkey/monkey_island/cc/services/__init__.py +++ b/monkey/monkey_island/cc/services/__init__.py @@ -1,3 +1,4 @@ from .authentication_service import AuthenticationService from .aws import AWSService +from .agent_signals_service import AgentSignalsService diff --git a/monkey/monkey_island/cc/services/agent_signals_service.py b/monkey/monkey_island/cc/services/agent_signals_service.py new file mode 100644 index 000000000..47c677271 --- /dev/null +++ b/monkey/monkey_island/cc/services/agent_signals_service.py @@ -0,0 +1,26 @@ +from datetime import datetime + +from monkey_island.cc.models import Agent, AgentSignals +from monkey_island.cc.repository import ISimulationRepository + + +class AgentSignalsService: + def __init__(self, simulation_repository: ISimulationRepository): + self._simulation_repository = simulation_repository + + def get_signals(self, agent: Agent) -> AgentSignals: + """ + Gets the signals sent to a particular agent + + :param agent: The agent whose signals need to be retrieved + :return: Signals sent to the relevant agent + """ + return AgentSignals(timestamp=datetime.now()) + + def on_terminate_agents_signal(self, timestamp: datetime): + """ + Updates the simulation repository with the terminate signal's timestamp + + :param timestamp: Timestamp of the terminate signal + """ + pass diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 742450c50..c63b26f3e 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -8,15 +8,14 @@ from common.agent_configuration.agent_sub_configurations import ( ScanTargetConfiguration, ) from common.credentials import Credentials, LMHash, NTHash -from common.types import SocketAddress from infection_monkey.exploit.log4shell_utils.ldap_server import LDAPServerFactory from monkey_island.cc.event_queue import IslandEventTopic, PyPubSubIslandEventQueue -from monkey_island.cc.models import AgentSignals, Report, Simulation +from monkey_island.cc.models import Report, Simulation +from monkey_island.cc.models import AgentSignals, Report from monkey_island.cc.models.networkmap import Arc, NetworkMap from monkey_island.cc.repository import MongoAgentRepository, MongoMachineRepository from monkey_island.cc.repository.attack.IMitigationsRepository import IMitigationsRepository from monkey_island.cc.repository.i_agent_event_repository import IAgentEventRepository -from monkey_island.cc.repository.i_agent_log_repository import IAgentLogRepository from monkey_island.cc.repository.i_agent_repository import IAgentRepository from monkey_island.cc.repository.i_attack_repository import IAttackRepository from monkey_island.cc.repository.i_config_repository import IConfigRepository @@ -27,6 +26,7 @@ from monkey_island.cc.repository.i_simulation_repository import ISimulationRepos from monkey_island.cc.repository.ICredentials import ICredentialsRepository from monkey_island.cc.repository.zero_trust.IEventRepository import IEventRepository from monkey_island.cc.repository.zero_trust.IFindingRepository import IFindingRepository +from monkey_island.cc.services import AgentSignalsService fake_monkey_dir_path # unused variable (monkey/tests/infection_monkey/post_breach/actions/test_users_custom_pba.py:37) set_os_linux # unused variable (monkey/tests/infection_monkey/post_breach/actions/test_users_custom_pba.py:37) @@ -302,11 +302,6 @@ IAgentEventRepository.get_events_by_type IAgentEventRepository.get_events_by_tag IAgentEventRepository.get_events_by_source -# TODO: Remove once #2274 is closed -IAgentLogRepository -IAgentLogRepository.upsert_agent_log -IAgentLogRepository.get_agent_log - # pydantic base models underscore_attrs_are_private @@ -326,7 +321,13 @@ EXPLOITED CC CC_TUNNEL -# TODO: Remove after #2323 -SocketAddress +IslandEventTopic.AGENT_CONNECTED +IslandEventTopic.CLEAR_SIMULATION_DATA +IslandEventTopic.RESET_AGENT_CONFIGURATION +# TODO: Remove after #2261 is closed +IslandEventTopic.TERMINATE_AGENTS Simulation.terminate_signal_time + +AgentSignalsService.get_signals +AgentSignalsService.on_terminate_agents_signal From 0775449fa95bbabc8d779f5b93778deef0c5ccfa Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Thu, 22 Sep 2022 17:42:43 +0000 Subject: [PATCH 05/67] Island: Add AgentSignals resource --- monkey/monkey_island/cc/app.py | 2 + monkey/monkey_island/cc/resources/__init__.py | 1 + .../cc/resources/agent_signals.py | 38 +++++++++ .../cc/resources/test_agent_signals.py | 83 +++++++++++++++++++ 4 files changed, 124 insertions(+) create mode 100644 monkey/monkey_island/cc/resources/agent_signals.py create mode 100644 monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index c56e13322..41d03372b 100644 --- a/monkey/monkey_island/cc/app.py +++ b/monkey/monkey_island/cc/app.py @@ -15,6 +15,7 @@ from monkey_island.cc.resources import ( AgentConfiguration, AgentEvents, Agents, + AgentSignals, ClearSimulationData, IPAddresses, IslandLog, @@ -188,6 +189,7 @@ def init_restful_endpoints(api: FlaskDIWrapper): api.add_resource(IPAddresses) api.add_resource(AgentEvents) + api.add_resource(AgentSignals) # API Spec: These two should be the same resource, GET for download and POST for upload api.add_resource(PBAFileDownload) diff --git a/monkey/monkey_island/cc/resources/__init__.py b/monkey/monkey_island/cc/resources/__init__.py index b13c6cc06..b99730e45 100644 --- a/monkey/monkey_island/cc/resources/__init__.py +++ b/monkey/monkey_island/cc/resources/__init__.py @@ -10,3 +10,4 @@ from .pba_file_upload import PBAFileUpload, LINUX_PBA_TYPE, WINDOWS_PBA_TYPE from .pba_file_download import PBAFileDownload from .agent_events import AgentEvents from .agents import Agents +from .agent_signals import AgentSignals diff --git a/monkey/monkey_island/cc/resources/agent_signals.py b/monkey/monkey_island/cc/resources/agent_signals.py new file mode 100644 index 000000000..8f3b800d4 --- /dev/null +++ b/monkey/monkey_island/cc/resources/agent_signals.py @@ -0,0 +1,38 @@ +import logging +from http import HTTPStatus +from json import JSONDecodeError + +from flask import request + +from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic +from monkey_island.cc.models import AgentSignals as Signal +from monkey_island.cc.resources.AbstractResource import AbstractResource + +logger = logging.getLogger(__name__) + + +class AgentSignals(AbstractResource): + urls = ["/api/agent-signals/terminate-all", "/api/agent-signals/"] + + def __init__( + self, + island_event_queue: IIslandEventQueue, + ): + self._island_event_queue = island_event_queue + + def post(self): + try: + signal = Signal(**request.json) + + # We allow an empty timestamp. However, should the agent be able to send us one? + if signal.terminate is None: + raise ValueError + self._island_event_queue.publish(IslandEventTopic.TERMINATE_AGENTS, signal=signal) + except (JSONDecodeError, TypeError, ValueError) as err: + return {"error": err}, HTTPStatus.BAD_REQUEST + + return {}, HTTPStatus.NO_CONTENT + + def get(self, agent_id: str): + # TODO: return AgentSignals + return {}, HTTPStatus.OK diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py new file mode 100644 index 000000000..2996bea93 --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py @@ -0,0 +1,83 @@ +from http import HTTPStatus +from unittest.mock import MagicMock +from uuid import UUID + +import pytest +from tests.common import StubDIContainer + +from monkey_island.cc.event_queue import IIslandEventQueue +from monkey_island.cc.resources import AgentSignals + +TIMESTAMP = 123456789 + + +@pytest.fixture( + params=[ + UUID("c0dd10b3-e21a-4da9-9d96-a99c19ebd7c5"), + UUID("9b4279f6-6ec5-4953-821e-893ddc71a988"), + ] +) +def agent_id(request) -> UUID: + return request.param + + +@pytest.fixture +def agent_signals_url(agent_id: UUID) -> str: + return f"/api/agent-signals/{agent_id}" + + +@pytest.fixture +def flask_client_builder(build_flask_client): + def inner(side_effect=None): + container = StubDIContainer() + + # TODO: Add AgentSignalsService and add values on publish + mock_island_event_queue = MagicMock(spec=IIslandEventQueue) + mock_island_event_queue.publish.side_effect = side_effect + container.register_instance(IIslandEventQueue, mock_island_event_queue) + + with build_flask_client(container) as flask_client: + return flask_client + + return inner + + +@pytest.fixture +def flask_client(flask_client_builder): + return flask_client_builder() + + +def test_agent_signals_terminate_all_post(flask_client): + resp = flask_client.post( + AgentSignals.urls[0], + json={"terminate": TIMESTAMP}, + follow_redirects=True, + ) + assert resp.status_code == HTTPStatus.NO_CONTENT + + +@pytest.mark.parametrize( + "bad_data", + [ + "bad timestamp", + {}, + {"wrong_key": TIMESTAMP}, + {"extra_key": "blah", "terminate": TIMESTAMP}, + TIMESTAMP, + ], +) +def test_agent_signals_terminate_all_post__invalid_timestamp(flask_client, bad_data): + resp = flask_client.post( + AgentSignals.urls[0], + json=bad_data, + follow_redirects=True, + ) + assert resp.status_code == HTTPStatus.BAD_REQUEST + + +# TODO: Complete this when GET is implemented +# Do we get a value indicating that we should stop? Depends on whether a signal was sent +def test_agent_signals_endpoint(flask_client, agent_signals_url): + resp = flask_client.get(agent_signals_url, follow_redirects=True) + assert resp.status_code == HTTPStatus.OK + assert resp.json == {} From dccef0efa5086887d84f68375274afca569b6389 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 13:29:58 +0530 Subject: [PATCH 06/67] Island: Rename Signal -> Signals in cc/resources/agent_signals.py --- monkey/monkey_island/cc/resources/agent_signals.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/resources/agent_signals.py b/monkey/monkey_island/cc/resources/agent_signals.py index 8f3b800d4..87e98fbee 100644 --- a/monkey/monkey_island/cc/resources/agent_signals.py +++ b/monkey/monkey_island/cc/resources/agent_signals.py @@ -5,7 +5,7 @@ from json import JSONDecodeError from flask import request from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic -from monkey_island.cc.models import AgentSignals as Signal +from monkey_island.cc.models import AgentSignals as Signals from monkey_island.cc.resources.AbstractResource import AbstractResource logger = logging.getLogger(__name__) @@ -22,7 +22,7 @@ class AgentSignals(AbstractResource): def post(self): try: - signal = Signal(**request.json) + signal = Signals(**request.json) # We allow an empty timestamp. However, should the agent be able to send us one? if signal.terminate is None: From 58ad44366a4297538daf2ade0c07a765d8885c97 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 13:30:55 +0530 Subject: [PATCH 07/67] Island: Remove comment in cc/resources/agent_signals.py --- monkey/monkey_island/cc/resources/agent_signals.py | 1 - 1 file changed, 1 deletion(-) diff --git a/monkey/monkey_island/cc/resources/agent_signals.py b/monkey/monkey_island/cc/resources/agent_signals.py index 87e98fbee..e04a9ea44 100644 --- a/monkey/monkey_island/cc/resources/agent_signals.py +++ b/monkey/monkey_island/cc/resources/agent_signals.py @@ -24,7 +24,6 @@ class AgentSignals(AbstractResource): try: signal = Signals(**request.json) - # We allow an empty timestamp. However, should the agent be able to send us one? if signal.terminate is None: raise ValueError self._island_event_queue.publish(IslandEventTopic.TERMINATE_AGENTS, signal=signal) From f23a6c8fa40982230d375874ac5e99e1c9416c0e Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 13:35:49 +0530 Subject: [PATCH 08/67] Island: Add message to ValueError in AgentSignals resource --- monkey/monkey_island/cc/resources/agent_signals.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/resources/agent_signals.py b/monkey/monkey_island/cc/resources/agent_signals.py index e04a9ea44..e55a489d4 100644 --- a/monkey/monkey_island/cc/resources/agent_signals.py +++ b/monkey/monkey_island/cc/resources/agent_signals.py @@ -25,7 +25,7 @@ class AgentSignals(AbstractResource): signal = Signals(**request.json) if signal.terminate is None: - raise ValueError + raise ValueError("Terminate signal's timestamp is empty") self._island_event_queue.publish(IslandEventTopic.TERMINATE_AGENTS, signal=signal) except (JSONDecodeError, TypeError, ValueError) as err: return {"error": err}, HTTPStatus.BAD_REQUEST From cfe31f8dee3c5d7323d160382769cf294e746eef Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 13:41:09 +0530 Subject: [PATCH 09/67] Island: Use terminate signal's timestamp directly instead of creating an AgentSignals object in AgentSignals resource --- monkey/monkey_island/cc/resources/agent_signals.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/monkey/monkey_island/cc/resources/agent_signals.py b/monkey/monkey_island/cc/resources/agent_signals.py index e55a489d4..b2abba820 100644 --- a/monkey/monkey_island/cc/resources/agent_signals.py +++ b/monkey/monkey_island/cc/resources/agent_signals.py @@ -5,7 +5,6 @@ from json import JSONDecodeError from flask import request from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic -from monkey_island.cc.models import AgentSignals as Signals from monkey_island.cc.resources.AbstractResource import AbstractResource logger = logging.getLogger(__name__) @@ -22,11 +21,14 @@ class AgentSignals(AbstractResource): def post(self): try: - signal = Signals(**request.json) - - if signal.terminate is None: + terminate_timestamp = request.json["kill_time"] + if terminate_timestamp is None: raise ValueError("Terminate signal's timestamp is empty") - self._island_event_queue.publish(IslandEventTopic.TERMINATE_AGENTS, signal=signal) + + self._island_event_queue.publish( + IslandEventTopic.TERMINATE_AGENTS, timestamp=terminate_timestamp + ) + except (JSONDecodeError, TypeError, ValueError) as err: return {"error": err}, HTTPStatus.BAD_REQUEST From cca4cf9df2a6e8ee5ba333590396e048dba541cc Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 13:52:43 +0530 Subject: [PATCH 10/67] Island: Implement AgentSignals resource's GET --- monkey/monkey_island/cc/resources/agent_signals.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/resources/agent_signals.py b/monkey/monkey_island/cc/resources/agent_signals.py index b2abba820..2b4439231 100644 --- a/monkey/monkey_island/cc/resources/agent_signals.py +++ b/monkey/monkey_island/cc/resources/agent_signals.py @@ -6,6 +6,7 @@ from flask import request from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.services import AgentSignalsService logger = logging.getLogger(__name__) @@ -16,8 +17,10 @@ class AgentSignals(AbstractResource): def __init__( self, island_event_queue: IIslandEventQueue, + agent_signals_service: AgentSignalsService, ): self._island_event_queue = island_event_queue + self._agent_signals_service = agent_signals_service def post(self): try: @@ -35,5 +38,5 @@ class AgentSignals(AbstractResource): return {}, HTTPStatus.NO_CONTENT def get(self, agent_id: str): - # TODO: return AgentSignals - return {}, HTTPStatus.OK + agent_signals = self._agent_signals_service.get_signals(agent_id) + return agent_signals.dict(), HTTPStatus.OK From 1afe625395c10456ca684cb7c6f208585fc4536e Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 14:02:50 +0530 Subject: [PATCH 11/67] Island: Catch KeyError in AgentSignals resource's POST --- monkey/monkey_island/cc/resources/agent_signals.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/resources/agent_signals.py b/monkey/monkey_island/cc/resources/agent_signals.py index 2b4439231..71afe4b71 100644 --- a/monkey/monkey_island/cc/resources/agent_signals.py +++ b/monkey/monkey_island/cc/resources/agent_signals.py @@ -32,7 +32,7 @@ class AgentSignals(AbstractResource): IslandEventTopic.TERMINATE_AGENTS, timestamp=terminate_timestamp ) - except (JSONDecodeError, TypeError, ValueError) as err: + except (JSONDecodeError, TypeError, ValueError, KeyError) as err: return {"error": err}, HTTPStatus.BAD_REQUEST return {}, HTTPStatus.NO_CONTENT From 5bf63c122145260986950b526d39f85bade23447 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 14:04:00 +0530 Subject: [PATCH 12/67] UT: Fix POST tests in test_agent_signals.py --- .../monkey_island/cc/resources/test_agent_signals.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py index 2996bea93..63ecef8f2 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py @@ -7,6 +7,7 @@ from tests.common import StubDIContainer from monkey_island.cc.event_queue import IIslandEventQueue from monkey_island.cc.resources import AgentSignals +from monkey_island.cc.services import AgentSignalsService TIMESTAMP = 123456789 @@ -31,11 +32,13 @@ def flask_client_builder(build_flask_client): def inner(side_effect=None): container = StubDIContainer() - # TODO: Add AgentSignalsService and add values on publish mock_island_event_queue = MagicMock(spec=IIslandEventQueue) mock_island_event_queue.publish.side_effect = side_effect container.register_instance(IIslandEventQueue, mock_island_event_queue) + mock_agent_signals_service = MagicMock(spec=AgentSignalsService) + container.register_instance(AgentSignalsService, mock_agent_signals_service) + with build_flask_client(container) as flask_client: return flask_client @@ -50,7 +53,7 @@ def flask_client(flask_client_builder): def test_agent_signals_terminate_all_post(flask_client): resp = flask_client.post( AgentSignals.urls[0], - json={"terminate": TIMESTAMP}, + json={"kill_time": TIMESTAMP}, follow_redirects=True, ) assert resp.status_code == HTTPStatus.NO_CONTENT @@ -62,7 +65,6 @@ def test_agent_signals_terminate_all_post(flask_client): "bad timestamp", {}, {"wrong_key": TIMESTAMP}, - {"extra_key": "blah", "terminate": TIMESTAMP}, TIMESTAMP, ], ) From bc43f81a112edc3203de0a43c661f5568e3d4e2f Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 18:05:35 +0530 Subject: [PATCH 13/67] UI: Use /api/agent-signals/terminate-all instead of /api/monkey-control/stop-all-agents --- monkey/monkey_island/cc/ui/src/components/pages/MapPage.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/ui/src/components/pages/MapPage.js b/monkey/monkey_island/cc/ui/src/components/pages/MapPage.js index 13db80e25..11a67456f 100644 --- a/monkey/monkey_island/cc/ui/src/components/pages/MapPage.js +++ b/monkey/monkey_island/cc/ui/src/components/pages/MapPage.js @@ -84,7 +84,7 @@ class MapPageComponent extends AuthComponent { } killAllMonkeys = () => { - this.authFetch('/api/monkey-control/stop-all-agents', + this.authFetch('/api/agent-signals/terminate-all', { method: 'POST', headers: {'Content-Type': 'application/json'}, From 263fff28f369f1c9417d7dceeb8c8fd1da60808a Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 11:50:26 +0530 Subject: [PATCH 14/67] BB: Use /api/agent-signals/terminate-all instead of /api/monkey-control/stop-all-agents --- envs/monkey_zoo/blackbox/island_client/monkey_island_client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 351de3c32..8ffa0978e 100644 --- a/envs/monkey_zoo/blackbox/island_client/monkey_island_client.py +++ b/envs/monkey_zoo/blackbox/island_client/monkey_island_client.py @@ -89,7 +89,7 @@ class MonkeyIslandClient(object): @avoid_race_condition def kill_all_monkeys(self): response = self.requests.post_json( - "api/monkey-control/stop-all-agents", json={"kill_time": time.time()} + "api/agent-signals/terminate-all", json={"kill_time": time.time()} ) if response.ok: LOGGER.info("Killing all monkeys after the test.") From 637978648a2b9e748cbbd7d23d1014c14884c822 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 11:53:14 +0530 Subject: [PATCH 15/67] Island: Remove StopAllAgents resource --- monkey/monkey_island/cc/app.py | 3 +-- .../cc/resources/agent_controls/__init__.py | 1 - .../agent_controls/stop_all_agents.py | 27 ------------------- 3 files changed, 1 insertion(+), 30 deletions(-) delete mode 100644 monkey/monkey_island/cc/resources/agent_controls/stop_all_agents.py diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index 41d03372b..3d127d985 100644 --- a/monkey/monkey_island/cc/app.py +++ b/monkey/monkey_island/cc/app.py @@ -26,7 +26,7 @@ from monkey_island.cc.resources import ( ResetAgentConfiguration, ) from monkey_island.cc.resources.AbstractResource import AbstractResource -from monkey_island.cc.resources.agent_controls import StopAgentCheck, StopAllAgents +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 +199,6 @@ def init_restful_endpoints(api: FlaskDIWrapper): api.add_resource(RemoteRun) api.add_resource(Version) api.add_resource(StopAgentCheck) - api.add_resource(StopAllAgents) # 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/resources/agent_controls/__init__.py b/monkey/monkey_island/cc/resources/agent_controls/__init__.py index 211696e4c..4bc6d5b48 100644 --- a/monkey/monkey_island/cc/resources/agent_controls/__init__.py +++ b/monkey/monkey_island/cc/resources/agent_controls/__init__.py @@ -1,2 +1 @@ -from .stop_all_agents import StopAllAgents from .stop_agent_check import StopAgentCheck diff --git a/monkey/monkey_island/cc/resources/agent_controls/stop_all_agents.py b/monkey/monkey_island/cc/resources/agent_controls/stop_all_agents.py deleted file mode 100644 index c3d719bd8..000000000 --- a/monkey/monkey_island/cc/resources/agent_controls/stop_all_agents.py +++ /dev/null @@ -1,27 +0,0 @@ -import json - -from flask import make_response, request - -from monkey_island.cc.resources.AbstractResource import AbstractResource -from monkey_island.cc.resources.request_authentication import jwt_required -from monkey_island.cc.resources.utils.semaphores import agent_killing_mutex -from monkey_island.cc.services.infection_lifecycle import set_stop_all, should_agent_die - - -class StopAllAgents(AbstractResource): - # API Spec: This is an action and there's no "resource"; RPC-style endpoint? - urls = ["/api/monkey-control/stop-all-agents"] - - @jwt_required - def post(self): - with agent_killing_mutex: - data = json.loads(request.data) - if data["kill_time"]: - set_stop_all(data["kill_time"]) - return make_response({}, 200) - else: - return make_response({}, 400) - - # API Spec: This is the exact same thing as what's in StopAgentCheck - def get(self, monkey_guid): - return {"stop_agent": should_agent_die(monkey_guid)} From ef273bc1cfe9f6939ad8a48b1fcc9f98cfa158a3 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 12:02:27 +0530 Subject: [PATCH 16/67] Island: Remove set_stop_all() --- .../monkey_island/cc/services/infection_lifecycle.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/monkey/monkey_island/cc/services/infection_lifecycle.py b/monkey/monkey_island/cc/services/infection_lifecycle.py index 937a3abeb..e1e3a831e 100644 --- a/monkey/monkey_island/cc/services/infection_lifecycle.py +++ b/monkey/monkey_island/cc/services/infection_lifecycle.py @@ -12,16 +12,6 @@ from monkey_island.cc.services.reporting.report_generation_synchronisation impor logger = logging.getLogger(__name__) -def set_stop_all(time: float): - # This will use Agent and Simulation repositories - for monkey in Monkey.objects(): - monkey.should_stop = True - monkey.save() - agent_controls = AgentControls.objects.first() - agent_controls.last_stop_all = time - agent_controls.save() - - 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) From 41951511d0bf9df61aba65479f572d8ac8dd16bf Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 12:19:53 +0200 Subject: [PATCH 17/67] Island: Add simplify=true when returning AgentSignals in endpoint --- monkey/monkey_island/cc/resources/agent_signals.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/resources/agent_signals.py b/monkey/monkey_island/cc/resources/agent_signals.py index 71afe4b71..765ea5d8d 100644 --- a/monkey/monkey_island/cc/resources/agent_signals.py +++ b/monkey/monkey_island/cc/resources/agent_signals.py @@ -39,4 +39,4 @@ class AgentSignals(AbstractResource): def get(self, agent_id: str): agent_signals = self._agent_signals_service.get_signals(agent_id) - return agent_signals.dict(), HTTPStatus.OK + return agent_signals.dict(simplify=True), HTTPStatus.OK From 2d42355e2c990cc602e945b0650f7bbdce52b535 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 12:21:36 +0200 Subject: [PATCH 18/67] UT: Add tests for GET AgentSignals endpoint --- .../cc/resources/test_agent_signals.py | 50 ++++++++++++------- 1 file changed, 33 insertions(+), 17 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py index 63ecef8f2..6fa371b09 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py @@ -6,29 +6,28 @@ import pytest from tests.common import StubDIContainer from monkey_island.cc.event_queue import IIslandEventQueue +from monkey_island.cc.models import AgentSignals as Signals +from monkey_island.cc.repository import RetrievalError, StorageError from monkey_island.cc.resources import AgentSignals from monkey_island.cc.services import AgentSignalsService TIMESTAMP = 123456789 +TIMESTAMP_1 = 123546789 +SIGNALS = Signals(terminate=TIMESTAMP) +SIGNALS_1 = Signals(terminate=TIMESTAMP_1) -@pytest.fixture( - params=[ - UUID("c0dd10b3-e21a-4da9-9d96-a99c19ebd7c5"), - UUID("9b4279f6-6ec5-4953-821e-893ddc71a988"), - ] -) -def agent_id(request) -> UUID: - return request.param +AGENT_ID = UUID("c0dd10b3-e21a-4da9-9d96-a99c19ebd7c5") +AGENT_ID_1 = UUID("9b4279f6-6ec5-4953-821e-893ddc71a988") @pytest.fixture -def agent_signals_url(agent_id: UUID) -> str: - return f"/api/agent-signals/{agent_id}" +def mock_agent_signals_service(): + return MagicMock(spec=AgentSignalsService) @pytest.fixture -def flask_client_builder(build_flask_client): +def flask_client_builder(build_flask_client, mock_agent_signals_service): def inner(side_effect=None): container = StubDIContainer() @@ -36,7 +35,6 @@ def flask_client_builder(build_flask_client): mock_island_event_queue.publish.side_effect = side_effect container.register_instance(IIslandEventQueue, mock_island_event_queue) - mock_agent_signals_service = MagicMock(spec=AgentSignalsService) container.register_instance(AgentSignalsService, mock_agent_signals_service) with build_flask_client(container) as flask_client: @@ -77,9 +75,27 @@ def test_agent_signals_terminate_all_post__invalid_timestamp(flask_client, bad_d assert resp.status_code == HTTPStatus.BAD_REQUEST -# TODO: Complete this when GET is implemented -# Do we get a value indicating that we should stop? Depends on whether a signal was sent -def test_agent_signals_endpoint(flask_client, agent_signals_url): - resp = flask_client.get(agent_signals_url, follow_redirects=True) +@pytest.mark.parametrize( + "url, signals", + [(f"/api/agent-signals/{AGENT_ID}", SIGNALS), (f"/api/agent-signals/{AGENT_ID_1}", SIGNALS_1)], +) +def test_agent_signals_get(flask_client, mock_agent_signals_service, url, signals): + mock_agent_signals_service.get_signals.return_value = signals + resp = flask_client.get(url, follow_redirects=True) assert resp.status_code == HTTPStatus.OK - assert resp.json == {} + assert resp.json == signals.dict(simplify=True) + + +@pytest.mark.parametrize( + "url, error", + [ + (f"/api/agent-signals/{AGENT_ID}", RetrievalError), + (f"/api/agent-signals/{AGENT_ID_1}", StorageError), + ], +) +def test_agent_signals_get__internal_server_error( + flask_client, mock_agent_signals_service, url, error +): + mock_agent_signals_service.get_signals.side_effect = error + resp = flask_client.get(url, follow_redirects=True) + assert resp.status_code == HTTPStatus.INTERNAL_SERVER_ERROR From 14c615e238120f3f8882530c67a4feda01b127d2 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 17:13:30 +0530 Subject: [PATCH 19/67] Island: Rename some variables in test_agent_signals.py --- .../cc/resources/test_agent_signals.py | 25 +++++++++++-------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py index 6fa371b09..91a13dd78 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py @@ -11,14 +11,14 @@ from monkey_island.cc.repository import RetrievalError, StorageError from monkey_island.cc.resources import AgentSignals from monkey_island.cc.services import AgentSignalsService -TIMESTAMP = 123456789 -TIMESTAMP_1 = 123546789 +TIMESTAMP_1 = 123456789 +TIMESTAMP_2 = 123546789 -SIGNALS = Signals(terminate=TIMESTAMP) SIGNALS_1 = Signals(terminate=TIMESTAMP_1) +SIGNALS_2 = Signals(terminate=TIMESTAMP_2) -AGENT_ID = UUID("c0dd10b3-e21a-4da9-9d96-a99c19ebd7c5") -AGENT_ID_1 = UUID("9b4279f6-6ec5-4953-821e-893ddc71a988") +AGENT_ID_1 = UUID("c0dd10b3-e21a-4da9-9d96-a99c19ebd7c5") +AGENT_ID_2 = UUID("9b4279f6-6ec5-4953-821e-893ddc71a988") @pytest.fixture @@ -51,7 +51,7 @@ def flask_client(flask_client_builder): def test_agent_signals_terminate_all_post(flask_client): resp = flask_client.post( AgentSignals.urls[0], - json={"kill_time": TIMESTAMP}, + json={"kill_time": TIMESTAMP_1}, follow_redirects=True, ) assert resp.status_code == HTTPStatus.NO_CONTENT @@ -62,8 +62,8 @@ def test_agent_signals_terminate_all_post(flask_client): [ "bad timestamp", {}, - {"wrong_key": TIMESTAMP}, - TIMESTAMP, + {"wrong_key": TIMESTAMP_1}, + TIMESTAMP_1, ], ) def test_agent_signals_terminate_all_post__invalid_timestamp(flask_client, bad_data): @@ -77,7 +77,10 @@ def test_agent_signals_terminate_all_post__invalid_timestamp(flask_client, bad_d @pytest.mark.parametrize( "url, signals", - [(f"/api/agent-signals/{AGENT_ID}", SIGNALS), (f"/api/agent-signals/{AGENT_ID_1}", SIGNALS_1)], + [ + (f"/api/agent-signals/{AGENT_ID_1}", SIGNALS_1), + (f"/api/agent-signals/{AGENT_ID_2}", SIGNALS_2), + ], ) def test_agent_signals_get(flask_client, mock_agent_signals_service, url, signals): mock_agent_signals_service.get_signals.return_value = signals @@ -89,8 +92,8 @@ def test_agent_signals_get(flask_client, mock_agent_signals_service, url, signal @pytest.mark.parametrize( "url, error", [ - (f"/api/agent-signals/{AGENT_ID}", RetrievalError), - (f"/api/agent-signals/{AGENT_ID_1}", StorageError), + (f"/api/agent-signals/{AGENT_ID_1}", RetrievalError), + (f"/api/agent-signals/{AGENT_ID_2}", StorageError), ], ) def test_agent_signals_get__internal_server_error( From fbfebc6167381c3aa955bbf922feff7a34bd9311 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 18:06:15 +0530 Subject: [PATCH 20/67] UI: 'kill_time' -> 'terminate_time' --- monkey/monkey_island/cc/ui/src/components/pages/MapPage.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/ui/src/components/pages/MapPage.js b/monkey/monkey_island/cc/ui/src/components/pages/MapPage.js index 11a67456f..d3bc5ad2f 100644 --- a/monkey/monkey_island/cc/ui/src/components/pages/MapPage.js +++ b/monkey/monkey_island/cc/ui/src/components/pages/MapPage.js @@ -89,7 +89,7 @@ class MapPageComponent extends AuthComponent { method: 'POST', headers: {'Content-Type': 'application/json'}, // Python uses floating point seconds, Date.now uses milliseconds, so convert - body: JSON.stringify({kill_time: Date.now() / 1000.0}) + body: JSON.stringify({terminate_time: Date.now() / 1000.0}) }) .then(res => res.json()) .then(() => {this.setState({killPressed: true})}); From 489ead31d2a7725de2f33293fbd04030f993a583 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 18:06:28 +0530 Subject: [PATCH 21/67] Island: 'kill_time' -> 'terminate_time' --- monkey/monkey_island/cc/resources/agent_signals.py | 2 +- monkey/monkey_island/cc/services/infection_lifecycle.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/monkey/monkey_island/cc/resources/agent_signals.py b/monkey/monkey_island/cc/resources/agent_signals.py index 765ea5d8d..01c8f0b7f 100644 --- a/monkey/monkey_island/cc/resources/agent_signals.py +++ b/monkey/monkey_island/cc/resources/agent_signals.py @@ -24,7 +24,7 @@ class AgentSignals(AbstractResource): def post(self): try: - terminate_timestamp = request.json["kill_time"] + terminate_timestamp = request.json["terminate_time"] if terminate_timestamp is None: raise ValueError("Terminate signal's timestamp is empty") diff --git a/monkey/monkey_island/cc/services/infection_lifecycle.py b/monkey/monkey_island/cc/services/infection_lifecycle.py index e1e3a831e..ef11ce5b1 100644 --- a/monkey/monkey_island/cc/services/infection_lifecycle.py +++ b/monkey/monkey_island/cc/services/infection_lifecycle.py @@ -27,14 +27,14 @@ def _should_agent_stop(monkey: Monkey) -> bool: def _is_monkey_killed_manually(monkey: Monkey) -> bool: - kill_timestamp = AgentControls.objects.first().last_stop_all - if kill_timestamp is None: + 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(kill_timestamp) >= int(launch_timestamp) + return int(terminate_timestamp) >= int(launch_timestamp) def get_completed_steps(): From 11f443e6414969b8083ae1abec479c5e65f3cc9d Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 17:15:00 +0530 Subject: [PATCH 22/67] UT: 'kill_time' -> 'terminate_time' --- .../unit_tests/monkey_island/cc/resources/test_agent_signals.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py index 91a13dd78..673ae6de7 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py @@ -51,7 +51,7 @@ def flask_client(flask_client_builder): def test_agent_signals_terminate_all_post(flask_client): resp = flask_client.post( AgentSignals.urls[0], - json={"kill_time": TIMESTAMP_1}, + json={"terminate_time": TIMESTAMP_1}, follow_redirects=True, ) assert resp.status_code == HTTPStatus.NO_CONTENT From 1632d8b3e9a2313884aa54e52dfd76b16b4724ab Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 17:15:24 +0530 Subject: [PATCH 23/67] BB: 'kill_time' -> 'terminate_time' --- envs/monkey_zoo/blackbox/island_client/monkey_island_client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 8ffa0978e..b3edac090 100644 --- a/envs/monkey_zoo/blackbox/island_client/monkey_island_client.py +++ b/envs/monkey_zoo/blackbox/island_client/monkey_island_client.py @@ -89,7 +89,7 @@ class MonkeyIslandClient(object): @avoid_race_condition def kill_all_monkeys(self): response = self.requests.post_json( - "api/agent-signals/terminate-all", json={"kill_time": time.time()} + "api/agent-signals/terminate-all", json={"terminate_time": time.time()} ) if response.ok: LOGGER.info("Killing all monkeys after the test.") From 24210d4f6f81e0985cbf069b1698fcf034a7363b Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 17:18:16 +0530 Subject: [PATCH 24/67] Island: Add check that terminate timestamp is > 0 in AgentSignals resource --- monkey/monkey_island/cc/resources/agent_signals.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/monkey/monkey_island/cc/resources/agent_signals.py b/monkey/monkey_island/cc/resources/agent_signals.py index 01c8f0b7f..e67478175 100644 --- a/monkey/monkey_island/cc/resources/agent_signals.py +++ b/monkey/monkey_island/cc/resources/agent_signals.py @@ -27,6 +27,8 @@ class AgentSignals(AbstractResource): terminate_timestamp = request.json["terminate_time"] if terminate_timestamp is None: raise ValueError("Terminate signal's timestamp is empty") + elif terminate_timestamp <= 0: + raise ValueError("Terminate signal's timestamp is not a positive integer") self._island_event_queue.publish( IslandEventTopic.TERMINATE_AGENTS, timestamp=terminate_timestamp From 28c3cf581fe20cb5a431a48f6b689d47cf7ad212 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 17:20:53 +0530 Subject: [PATCH 25/67] UT: Add test cases for AgentSignal resource's POST --- .../unit_tests/monkey_island/cc/resources/test_agent_signals.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py index 673ae6de7..b4a6a4d7f 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py @@ -64,6 +64,8 @@ def test_agent_signals_terminate_all_post(flask_client): {}, {"wrong_key": TIMESTAMP_1}, TIMESTAMP_1, + {"terminate_time": 0}, + {"terminate_time": -1}, ], ) def test_agent_signals_terminate_all_post__invalid_timestamp(flask_client, bad_data): From 2864286a2956d97c1e2bbfff33252d702a9f0760 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 17:46:55 +0530 Subject: [PATCH 26/67] Island: Add TerminateAllAgents resource --- monkey/monkey_island/cc/app.py | 2 + .../cc/resources/terminate_all_agents.py | 39 +++++++++++++++++++ 2 files changed, 41 insertions(+) create mode 100644 monkey/monkey_island/cc/resources/terminate_all_agents.py diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index 3d127d985..7438f1eac 100644 --- a/monkey/monkey_island/cc/app.py +++ b/monkey/monkey_island/cc/app.py @@ -49,6 +49,7 @@ from monkey_island.cc.resources.root import Root from monkey_island.cc.resources.security_report import SecurityReport from monkey_island.cc.resources.telemetry import Telemetry from monkey_island.cc.resources.telemetry_feed import TelemetryFeed +from monkey_island.cc.resources.terminate_all_agents import TerminateAllAgents from monkey_island.cc.resources.version import Version from monkey_island.cc.resources.zero_trust.finding_event import ZeroTrustFindingEvent from monkey_island.cc.resources.zero_trust.zero_trust_report import ZeroTrustReport @@ -212,6 +213,7 @@ def init_restful_endpoints(api: FlaskDIWrapper): def init_rpc_endpoints(api: FlaskDIWrapper): api.add_resource(ResetAgentConfiguration) api.add_resource(ClearSimulationData) + api.add_resource(TerminateAllAgents) def init_app(mongo_url: str, container: DIContainer): diff --git a/monkey/monkey_island/cc/resources/terminate_all_agents.py b/monkey/monkey_island/cc/resources/terminate_all_agents.py new file mode 100644 index 000000000..70603305f --- /dev/null +++ b/monkey/monkey_island/cc/resources/terminate_all_agents.py @@ -0,0 +1,39 @@ +import logging +from http import HTTPStatus +from json import JSONDecodeError + +from flask import request + +from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required + +logger = logging.getLogger(__name__) + + +class TerminateAllAgents(AbstractResource): + urls = ["/api/terminate-all-agents"] + + def __init__( + self, + island_event_queue: IIslandEventQueue, + ): + self._island_event_queue = island_event_queue + + @jwt_required + def post(self): + try: + terminate_timestamp = request.json["terminate_time"] + if terminate_timestamp is None: + raise ValueError("Terminate signal's timestamp is empty") + elif terminate_timestamp <= 0: + raise ValueError("Terminate signal's timestamp is not a positive integer") + + self._island_event_queue.publish( + IslandEventTopic.TERMINATE_AGENTS, timestamp=terminate_timestamp + ) + + except (JSONDecodeError, TypeError, ValueError, KeyError) as err: + return {"error": err}, HTTPStatus.BAD_REQUEST + + return {}, HTTPStatus.NO_CONTENT From 7527eca86124bb5d67530c201c542a1bf3e21a4f Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 17:52:37 +0530 Subject: [PATCH 27/67] UI: '/api/agent-signals/terminate-all' -> '/api/terminate-all-agents' --- monkey/monkey_island/cc/ui/src/components/pages/MapPage.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/ui/src/components/pages/MapPage.js b/monkey/monkey_island/cc/ui/src/components/pages/MapPage.js index d3bc5ad2f..764243f60 100644 --- a/monkey/monkey_island/cc/ui/src/components/pages/MapPage.js +++ b/monkey/monkey_island/cc/ui/src/components/pages/MapPage.js @@ -84,7 +84,7 @@ class MapPageComponent extends AuthComponent { } killAllMonkeys = () => { - this.authFetch('/api/agent-signals/terminate-all', + this.authFetch('/api/terminate-all-agents', { method: 'POST', headers: {'Content-Type': 'application/json'}, From c586623b8b74b84343e3c185e4d67742de8650f3 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 17:52:54 +0530 Subject: [PATCH 28/67] BB: '/api/agent-signals/terminate-all' -> '/api/terminate-all-agents' --- envs/monkey_zoo/blackbox/island_client/monkey_island_client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 b3edac090..7b222555c 100644 --- a/envs/monkey_zoo/blackbox/island_client/monkey_island_client.py +++ b/envs/monkey_zoo/blackbox/island_client/monkey_island_client.py @@ -89,7 +89,7 @@ class MonkeyIslandClient(object): @avoid_race_condition def kill_all_monkeys(self): response = self.requests.post_json( - "api/agent-signals/terminate-all", json={"terminate_time": time.time()} + "api/terminate-all-agents", json={"terminate_time": time.time()} ) if response.ok: LOGGER.info("Killing all monkeys after the test.") From 105cc60f4b988c9e6d2af7da68504f324a8e85c9 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 17:54:00 +0530 Subject: [PATCH 29/67] Island: Remove POST method from AgentSignals resource --- .../cc/resources/agent_signals.py | 25 +------------------ 1 file changed, 1 insertion(+), 24 deletions(-) diff --git a/monkey/monkey_island/cc/resources/agent_signals.py b/monkey/monkey_island/cc/resources/agent_signals.py index e67478175..f6ce22f71 100644 --- a/monkey/monkey_island/cc/resources/agent_signals.py +++ b/monkey/monkey_island/cc/resources/agent_signals.py @@ -1,10 +1,6 @@ import logging from http import HTTPStatus -from json import JSONDecodeError -from flask import request - -from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic from monkey_island.cc.resources.AbstractResource import AbstractResource from monkey_island.cc.services import AgentSignalsService @@ -12,33 +8,14 @@ logger = logging.getLogger(__name__) class AgentSignals(AbstractResource): - urls = ["/api/agent-signals/terminate-all", "/api/agent-signals/"] + urls = ["/api/agent-signals/"] def __init__( self, - island_event_queue: IIslandEventQueue, agent_signals_service: AgentSignalsService, ): - self._island_event_queue = island_event_queue self._agent_signals_service = agent_signals_service - def post(self): - try: - terminate_timestamp = request.json["terminate_time"] - if terminate_timestamp is None: - raise ValueError("Terminate signal's timestamp is empty") - elif terminate_timestamp <= 0: - raise ValueError("Terminate signal's timestamp is not a positive integer") - - self._island_event_queue.publish( - IslandEventTopic.TERMINATE_AGENTS, timestamp=terminate_timestamp - ) - - except (JSONDecodeError, TypeError, ValueError, KeyError) as err: - return {"error": err}, HTTPStatus.BAD_REQUEST - - return {}, HTTPStatus.NO_CONTENT - def get(self, agent_id: str): agent_signals = self._agent_signals_service.get_signals(agent_id) return agent_signals.dict(simplify=True), HTTPStatus.OK From 645e03e46f896f1333c5124cd97c2a5c99dfc434 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 17:59:31 +0530 Subject: [PATCH 30/67] Island: Import TerminateAllAgents in cc/resources/__init__.py --- monkey/monkey_island/cc/resources/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/monkey/monkey_island/cc/resources/__init__.py b/monkey/monkey_island/cc/resources/__init__.py index b99730e45..bef3988e8 100644 --- a/monkey/monkey_island/cc/resources/__init__.py +++ b/monkey/monkey_island/cc/resources/__init__.py @@ -11,3 +11,4 @@ from .pba_file_download import PBAFileDownload from .agent_events import AgentEvents from .agents import Agents from .agent_signals import AgentSignals +from .terminate_all_agents import TerminateAllAgents From 066f106882cd95186791f58d4d4561beb655bb87 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 18:01:12 +0530 Subject: [PATCH 31/67] UT: Move relevant tests from test_agent_signals.py to test_terminate_all_agents.py --- .../cc/resources/test_agent_signals.py | 36 ----------- .../cc/resources/test_terminate_all_agents.py | 59 +++++++++++++++++++ 2 files changed, 59 insertions(+), 36 deletions(-) create mode 100644 monkey/tests/unit_tests/monkey_island/cc/resources/test_terminate_all_agents.py diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py index b4a6a4d7f..3f644e051 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py @@ -5,10 +5,8 @@ from uuid import UUID import pytest from tests.common import StubDIContainer -from monkey_island.cc.event_queue import IIslandEventQueue from monkey_island.cc.models import AgentSignals as Signals from monkey_island.cc.repository import RetrievalError, StorageError -from monkey_island.cc.resources import AgentSignals from monkey_island.cc.services import AgentSignalsService TIMESTAMP_1 = 123456789 @@ -30,11 +28,6 @@ def mock_agent_signals_service(): def flask_client_builder(build_flask_client, mock_agent_signals_service): def inner(side_effect=None): container = StubDIContainer() - - mock_island_event_queue = MagicMock(spec=IIslandEventQueue) - mock_island_event_queue.publish.side_effect = side_effect - container.register_instance(IIslandEventQueue, mock_island_event_queue) - container.register_instance(AgentSignalsService, mock_agent_signals_service) with build_flask_client(container) as flask_client: @@ -48,35 +41,6 @@ def flask_client(flask_client_builder): return flask_client_builder() -def test_agent_signals_terminate_all_post(flask_client): - resp = flask_client.post( - AgentSignals.urls[0], - json={"terminate_time": TIMESTAMP_1}, - follow_redirects=True, - ) - assert resp.status_code == HTTPStatus.NO_CONTENT - - -@pytest.mark.parametrize( - "bad_data", - [ - "bad timestamp", - {}, - {"wrong_key": TIMESTAMP_1}, - TIMESTAMP_1, - {"terminate_time": 0}, - {"terminate_time": -1}, - ], -) -def test_agent_signals_terminate_all_post__invalid_timestamp(flask_client, bad_data): - resp = flask_client.post( - AgentSignals.urls[0], - json=bad_data, - follow_redirects=True, - ) - assert resp.status_code == HTTPStatus.BAD_REQUEST - - @pytest.mark.parametrize( "url, signals", [ diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_terminate_all_agents.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_terminate_all_agents.py new file mode 100644 index 000000000..305293f26 --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_terminate_all_agents.py @@ -0,0 +1,59 @@ +from http import HTTPStatus +from unittest.mock import MagicMock + +import pytest +from tests.common import StubDIContainer + +from monkey_island.cc.event_queue import IIslandEventQueue +from monkey_island.cc.resources import TerminateAllAgents + +TIMESTAMP = 123456789 + + +@pytest.fixture +def flask_client_builder(build_flask_client): + def inner(side_effect=None): + container = StubDIContainer() + + mock_island_event_queue = MagicMock(spec=IIslandEventQueue) + mock_island_event_queue.publish.side_effect = side_effect + container.register_instance(IIslandEventQueue, mock_island_event_queue) + + with build_flask_client(container) as flask_client: + return flask_client + + return inner + + +@pytest.fixture +def flask_client(flask_client_builder): + return flask_client_builder() + + +def test_terminate_all_agents_post(flask_client): + resp = flask_client.post( + TerminateAllAgents.urls[0], + json={"terminate_time": TIMESTAMP}, + follow_redirects=True, + ) + assert resp.status_code == HTTPStatus.NO_CONTENT + + +@pytest.mark.parametrize( + "bad_data", + [ + "bad timestamp", + {}, + {"wrong_key": TIMESTAMP}, + TIMESTAMP, + {"terminate_time": 0}, + {"terminate_time": -1}, + ], +) +def test_terminate_all_agents_post__invalid_timestamp(flask_client, bad_data): + resp = flask_client.post( + TerminateAllAgents.urls[0], + json=bad_data, + follow_redirects=True, + ) + assert resp.status_code == HTTPStatus.BAD_REQUEST From d10c148533e15d00ed37a3c83d8b3e338c706274 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Thu, 22 Sep 2022 14:43:04 +0200 Subject: [PATCH 32/67] Island: Add `get_progenitor` to IAgentRepository --- .../monkey_island/cc/repository/i_agent_repository.py | 11 +++++++++++ vulture_allowlist.py | 2 ++ 2 files changed, 13 insertions(+) diff --git a/monkey/monkey_island/cc/repository/i_agent_repository.py b/monkey/monkey_island/cc/repository/i_agent_repository.py index 880b4d7e7..ae12e8eea 100644 --- a/monkey/monkey_island/cc/repository/i_agent_repository.py +++ b/monkey/monkey_island/cc/repository/i_agent_repository.py @@ -40,6 +40,17 @@ class IAgentRepository(ABC): :raises RetrievalError: If an error occurred while attempting to retrieve the `Agents` """ + @abstractmethod + def get_progenitor(self, agent: Agent) -> Agent: + """ + Gets the progenitor `Agent` for the agent. + + :param agent: The Agent for which we want the progenitor + :return: `Agent` progenitor ( an initial agent that started the exploitation chain) + :raises RetrievalError: If an error occured while attempting to retrieve the `Agent` + :raises UnknownRecordError: If the agent ID is not in the repository + """ + @abstractmethod def reset(self): """ diff --git a/vulture_allowlist.py b/vulture_allowlist.py index ef9613fca..1e4196b78 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -246,6 +246,8 @@ IMitigationsRepository.save_mitigations IAgentRepository.upsert_agent IAgentRepository.get_agent_by_id IAgentRepository.get_running_agents +IAgentRepository.get_progenitor +descendant agent IAttackRepository.get_attack_report IAttackRepository.save_attack_report From b666078e7d1743b5be9faef012fc553c6a31141a Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 18:31:37 +0530 Subject: [PATCH 33/67] Island: '/api/terminate-all-agents' -> '/api/agent-signals/terminate-all-agents' --- monkey/monkey_island/cc/resources/terminate_all_agents.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/resources/terminate_all_agents.py b/monkey/monkey_island/cc/resources/terminate_all_agents.py index 70603305f..70032f11e 100644 --- a/monkey/monkey_island/cc/resources/terminate_all_agents.py +++ b/monkey/monkey_island/cc/resources/terminate_all_agents.py @@ -12,7 +12,7 @@ logger = logging.getLogger(__name__) class TerminateAllAgents(AbstractResource): - urls = ["/api/terminate-all-agents"] + urls = ["/api/agent-signals/terminate-all-agents"] def __init__( self, From f12e8398787b8cf2f0078a7f688f876c08ac032e Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Thu, 22 Sep 2022 17:34:57 +0200 Subject: [PATCH 34/67] Island: Register AgentSignalsSerivce in DI Container --- monkey/monkey_island/cc/services/initialize.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index 830e010d7..c0dc600c3 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -47,7 +47,7 @@ from monkey_island.cc.repository import ( ) from monkey_island.cc.server_utils.consts import MONKEY_ISLAND_ABS_PATH from monkey_island.cc.server_utils.encryption import ILockableEncryptor, RepositoryEncryptor -from monkey_island.cc.services import AWSService +from monkey_island.cc.services import AgentSignalsService, AWSService from monkey_island.cc.services.attack.technique_reports.T1003 import T1003, T1003GetReportData from monkey_island.cc.services.run_local_monkey import LocalMonkeyRunService from monkey_island.cc.setup.mongo.mongo_setup import MONGO_URL @@ -176,6 +176,7 @@ def _register_services(container: DIContainer): container.register_instance(AWSService, container.resolve(AWSService)) container.register_instance(LocalMonkeyRunService, container.resolve(LocalMonkeyRunService)) container.register_instance(AuthenticationService, container.resolve(AuthenticationService)) + container.register_instance(AgentSignalsService, container.resolve(AgentSignalsService)) def _dirty_hacks(container: DIContainer): From c25e245a8e5b59400b8d1b616baf1556ff13e35b Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Thu, 22 Sep 2022 14:45:07 +0200 Subject: [PATCH 35/67] Island: Implement `get_progenitor` in MongoAgentRepository --- .../cc/repository/mongo_agent_repository.py | 8 ++++ .../repository/test_mongo_agent_repository.py | 37 ++++++++++++++++++- 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/repository/mongo_agent_repository.py b/monkey/monkey_island/cc/repository/mongo_agent_repository.py index dfad2bbf7..021cf34bc 100644 --- a/monkey/monkey_island/cc/repository/mongo_agent_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_agent_repository.py @@ -58,6 +58,14 @@ class MongoAgentRepository(IAgentRepository): except Exception as err: raise RetrievalError(f"Error retrieving running agents: {err}") + def get_progenitor(self, agent: Agent) -> Agent: + if agent.parent_id is None: + return agent + + parent = self.get_agent_by_id(agent.parent_id) + + return self.get_progenitor(parent) + def reset(self): try: self._agents_collection.drop() diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_repository.py index a1221d32d..938ea7412 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_repository.py @@ -17,14 +17,29 @@ from monkey_island.cc.repository import ( ) VICTIM_ZERO_ID = uuid4() +VICTIM_TWO_ID = uuid4() +VICTIM_THREE_ID = uuid4() + +PROGENITOR_AGENT = Agent( + id=VICTIM_ZERO_ID, machine_id=1, start_time=datetime.fromtimestamp(1661856718) +) + +DESCENDANT_AGENT = Agent( + id=VICTIM_THREE_ID, + machine_id=4, + start_time=datetime.fromtimestamp(1661856868), + parent_id=VICTIM_TWO_ID, +) + RUNNING_AGENTS = ( - Agent(id=VICTIM_ZERO_ID, machine_id=1, start_time=datetime.fromtimestamp(1661856718)), + PROGENITOR_AGENT, Agent( - id=uuid4(), + id=VICTIM_TWO_ID, machine_id=2, start_time=datetime.fromtimestamp(1661856818), parent_id=VICTIM_ZERO_ID, ), + DESCENDANT_AGENT, ) STOPPED_AGENTS = ( Agent( @@ -172,6 +187,24 @@ def test_get_running_agents__retrieval_error(error_raising_agent_repository): error_raising_agent_repository.get_running_agents() +@pytest.mark.parametrize("agent", [DESCENDANT_AGENT, PROGENITOR_AGENT]) +def test_get_progenitor(agent_repository, agent): + actual_progenitor = agent_repository.get_progenitor(agent) + + assert actual_progenitor == PROGENITOR_AGENT + + +def test_get_progenitor__id_not_found(agent_repository): + dummy_agent = Agent(id=uuid4(), machine_id=10, start_time=datetime.now(), parent_id=uuid4()) + with pytest.raises(UnknownRecordError): + agent_repository.get_progenitor(dummy_agent) + + +def test_get_progenitor__retrieval_error(error_raising_agent_repository): + with pytest.raises(RetrievalError): + error_raising_agent_repository.get_progenitor(AGENTS[1]) + + def test_reset(agent_repository): # Ensure the repository is not empty for agent in AGENTS: From 850857c8a1bf3ef972789b522d91a918d0b30206 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 18:32:18 +0530 Subject: [PATCH 36/67] UI: '/api/terminate-all-agents' -> '/api/agent-signals/terminate-all-agents' --- monkey/monkey_island/cc/ui/src/components/pages/MapPage.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/ui/src/components/pages/MapPage.js b/monkey/monkey_island/cc/ui/src/components/pages/MapPage.js index 764243f60..9b6a76eb4 100644 --- a/monkey/monkey_island/cc/ui/src/components/pages/MapPage.js +++ b/monkey/monkey_island/cc/ui/src/components/pages/MapPage.js @@ -84,7 +84,7 @@ class MapPageComponent extends AuthComponent { } killAllMonkeys = () => { - this.authFetch('/api/terminate-all-agents', + this.authFetch('/api/agent-signals/terminate-all-agents', { method: 'POST', headers: {'Content-Type': 'application/json'}, From 5eeee2a60d7bf2873f02bb2b47d38e3f559b58a9 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Thu, 22 Sep 2022 17:36:42 +0200 Subject: [PATCH 37/67] Island: Subscribe AgentSignalsService.on_terminate_agents_signal to TERMINATE_AGENTS events --- .../monkey_island/cc/setup/island_event_handlers.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/monkey/monkey_island/cc/setup/island_event_handlers.py b/monkey/monkey_island/cc/setup/island_event_handlers.py index 6b4cb8c53..31e7908f6 100644 --- a/monkey/monkey_island/cc/setup/island_event_handlers.py +++ b/monkey/monkey_island/cc/setup/island_event_handlers.py @@ -15,6 +15,7 @@ from monkey_island.cc.repository import ( INodeRepository, ISimulationRepository, ) +from monkey_island.cc.services import AgentSignalsService from monkey_island.cc.services.database import Database @@ -25,6 +26,7 @@ def setup_island_event_handlers(container: DIContainer): _subscribe_reset_agent_configuration_events(island_event_queue, container) _subscribe_clear_simulation_data_events(island_event_queue, container) _subscribe_set_island_mode_events(island_event_queue, container) + _subscribe_on_terminate_agents_signal(island_event_queue, container) def _subscribe_agent_registration_events( @@ -74,3 +76,13 @@ def _subscribe_set_island_mode_events( simulation_repository = container.resolve(ISimulationRepository) island_event_queue.subscribe(topic, simulation_repository.set_mode) + + +def _subscribe_on_terminate_agents_signal( + island_event_queue: IIslandEventQueue, container: DIContainer +): + topic = IslandEventTopic.TERMINATE_AGENTS + + agent_signals_service = container.resolve(AgentSignalsService) + + island_event_queue.subscribe(topic, agent_signals_service.on_terminate_agents_signal) From 6174e8dfcbdfbf4d97f1dd0149a4dd3dadaf8894 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 18:32:27 +0530 Subject: [PATCH 38/67] BB: '/api/terminate-all-agents' -> '/api/agent-signals/terminate-all-agents' --- envs/monkey_zoo/blackbox/island_client/monkey_island_client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 7b222555c..f1df4a25b 100644 --- a/envs/monkey_zoo/blackbox/island_client/monkey_island_client.py +++ b/envs/monkey_zoo/blackbox/island_client/monkey_island_client.py @@ -89,7 +89,7 @@ class MonkeyIslandClient(object): @avoid_race_condition def kill_all_monkeys(self): response = self.requests.post_json( - "api/terminate-all-agents", json={"terminate_time": time.time()} + "api/agent-signals/terminate-all-agents", json={"terminate_time": time.time()} ) if response.ok: LOGGER.info("Killing all monkeys after the test.") From 9d3be7e1d3936fff758d94fa8e05e3bfa15112e2 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Thu, 22 Sep 2022 17:37:14 +0200 Subject: [PATCH 39/67] Island: Implement AgentSignalsService.on_terminate_agents_signal --- .../monkey_island/cc/services/agent_signals_service.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/services/agent_signals_service.py b/monkey/monkey_island/cc/services/agent_signals_service.py index 23a15dba9..473a10066 100644 --- a/monkey/monkey_island/cc/services/agent_signals_service.py +++ b/monkey/monkey_island/cc/services/agent_signals_service.py @@ -1,10 +1,13 @@ +import logging from datetime import datetime from typing import Optional from common.types import AgentID -from monkey_island.cc.models import AgentSignals +from monkey_island.cc.models import AgentSignals, Simulation from monkey_island.cc.repository import IAgentRepository, ISimulationRepository +logger = logging.getLogger(__name__) + class AgentSignalsService: def __init__( @@ -45,4 +48,7 @@ class AgentSignalsService: :param timestamp: Timestamp of the terminate signal """ - pass + simulation = self._simulation_repository.get_simulation() + updated_simulation = Simulation(mode=simulation.mode, terminate_signal_time=timestamp) + + self._simulation_repository.save_simulation(updated_simulation) From f7997a6a509f4db335bfd93a3fc7771ede63e0a7 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 22 Sep 2022 10:04:56 -0400 Subject: [PATCH 40/67] Island: Fix tenses in repository docstrings --- .../monkey_island/cc/repository/i_agent_repository.py | 8 ++++---- .../cc/repository/i_machine_repository.py | 10 +++++----- .../monkey_island/cc/repository/i_node_repository.py | 6 +++--- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/monkey/monkey_island/cc/repository/i_agent_repository.py b/monkey/monkey_island/cc/repository/i_agent_repository.py index ae12e8eea..619738ab5 100644 --- a/monkey/monkey_island/cc/repository/i_agent_repository.py +++ b/monkey/monkey_island/cc/repository/i_agent_repository.py @@ -16,7 +16,7 @@ class IAgentRepository(ABC): already exists, update it. :param agent: The `agent` to be inserted or updated - :raises StorageError: If an error occurred while attempting to store the `Agent` + :raises StorageError: If an error occurs while attempting to store the `Agent` """ @abstractmethod @@ -28,7 +28,7 @@ class IAgentRepository(ABC): :return: An `Agent` with a matching `id` :raises UnknownRecordError: If an `Agent` with the specified `id` does not exist in the repository - :raises RetrievalError: If an error occurred while attempting to retrieve the `Agent` + :raises RetrievalError: If an error occurs while attempting to retrieve the `Agent` """ @abstractmethod @@ -37,7 +37,7 @@ class IAgentRepository(ABC): Get all `Agents` that are currently running :return: All `Agents` that are currently running - :raises RetrievalError: If an error occurred while attempting to retrieve the `Agents` + :raises RetrievalError: If an error occurs while attempting to retrieve the `Agents` """ @abstractmethod @@ -47,7 +47,7 @@ class IAgentRepository(ABC): :param agent: The Agent for which we want the progenitor :return: `Agent` progenitor ( an initial agent that started the exploitation chain) - :raises RetrievalError: If an error occured while attempting to retrieve the `Agent` + :raises RetrievalError: If an error occurrs while attempting to retrieve the `Agent` :raises UnknownRecordError: If the agent ID is not in the repository """ diff --git a/monkey/monkey_island/cc/repository/i_machine_repository.py b/monkey/monkey_island/cc/repository/i_machine_repository.py index d60346eb2..85e8e71f3 100644 --- a/monkey/monkey_island/cc/repository/i_machine_repository.py +++ b/monkey/monkey_island/cc/repository/i_machine_repository.py @@ -26,7 +26,7 @@ class IMachineRepository(ABC): `Machine` already exists, update it. :param machine: The `Machine` to be inserted or updated - :raises StorageError: If an error occurred while attempting to store the `Machine` + :raises StorageError: If an error occurs while attempting to store the `Machine` """ @abstractmethod @@ -38,7 +38,7 @@ class IMachineRepository(ABC): :return: A `Machine` with a matching `id` :raises UnknownRecordError: If a `Machine` with the specified `id` does not exist in the repository - :raises RetrievalError: If an error occurred while attempting to retrieve the `Machine` + :raises RetrievalError: If an error occurs while attempting to retrieve the `Machine` """ @abstractmethod @@ -50,7 +50,7 @@ class IMachineRepository(ABC): :return: A `Machine` with a matching `hardware_id` :raises UnknownRecordError: If a `Machine` with the specified `hardware_id` does not exist in the repository - :raises RetrievalError: If an error occurred while attempting to retrieve the `Machine` + :raises RetrievalError: If an error occurs while attempting to retrieve the `Machine` """ @abstractmethod @@ -62,7 +62,7 @@ class IMachineRepository(ABC): :return: A sequence of Machines that have a network interface with a matching IP :raises UnknownRecordError: If a `Machine` with the specified `ip` does not exist in the repository - :raises RetrievalError: If an error occurred while attempting to retrieve the `Machine` + :raises RetrievalError: If an error occurs while attempting to retrieve the `Machine` """ @abstractmethod @@ -70,6 +70,6 @@ class IMachineRepository(ABC): """ Removes all data from the repository - :raises RemovalError: If an error occurred while attempting to remove all `Machines` from + :raises RemovalError: If an error occurs while attempting to remove all `Machines` from the repository """ diff --git a/monkey/monkey_island/cc/repository/i_node_repository.py b/monkey/monkey_island/cc/repository/i_node_repository.py index d848d58fe..53fc53bc9 100644 --- a/monkey/monkey_island/cc/repository/i_node_repository.py +++ b/monkey/monkey_island/cc/repository/i_node_repository.py @@ -22,7 +22,7 @@ class INodeRepository(ABC): :param src: The machine that the connection or communication originated from :param dst: The machine that the src communicated with :param communication_type: The way the machines communicated - :raises StorageError: If an error occurred while attempting to upsert the Node + :raises StorageError: If an error occurs while attempting to upsert the Node """ @abstractmethod @@ -31,7 +31,7 @@ class INodeRepository(ABC): Return all nodes that are stored in the repository :return: All known Nodes - :raises RetrievalError: If an error occurred while attempting to retrieve the nodes + :raises RetrievalError: If an error occurs while attempting to retrieve the nodes """ @abstractmethod @@ -39,6 +39,6 @@ class INodeRepository(ABC): """ Removes all data from the repository - :raises RemovalError: If an error occurred while attempting to remove all `Nodes` from the + :raises RemovalError: If an error occurs while attempting to remove all `Nodes` from the repository """ From c4642141f0358e95a354300ac8414c1d565e038e Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 18:36:03 +0530 Subject: [PATCH 41/67] Island: Move terminate_all_agents.py and terminate_all_agents.py under cc/resources/agent_signals/ --- monkey/monkey_island/cc/app.py | 2 +- monkey/monkey_island/cc/resources/__init__.py | 3 +-- monkey/monkey_island/cc/resources/agent_signals/__init__.py | 2 ++ .../cc/resources/{ => agent_signals}/agent_signals.py | 0 .../cc/resources/{ => agent_signals}/terminate_all_agents.py | 0 5 files changed, 4 insertions(+), 3 deletions(-) create mode 100644 monkey/monkey_island/cc/resources/agent_signals/__init__.py rename monkey/monkey_island/cc/resources/{ => agent_signals}/agent_signals.py (100%) rename monkey/monkey_island/cc/resources/{ => agent_signals}/terminate_all_agents.py (100%) diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index 7438f1eac..86616b596 100644 --- a/monkey/monkey_island/cc/app.py +++ b/monkey/monkey_island/cc/app.py @@ -24,6 +24,7 @@ from monkey_island.cc.resources import ( PropagationCredentials, RemoteRun, ResetAgentConfiguration, + TerminateAllAgents, ) from monkey_island.cc.resources.AbstractResource import AbstractResource from monkey_island.cc.resources.agent_controls import StopAgentCheck @@ -49,7 +50,6 @@ from monkey_island.cc.resources.root import Root from monkey_island.cc.resources.security_report import SecurityReport from monkey_island.cc.resources.telemetry import Telemetry from monkey_island.cc.resources.telemetry_feed import TelemetryFeed -from monkey_island.cc.resources.terminate_all_agents import TerminateAllAgents from monkey_island.cc.resources.version import Version from monkey_island.cc.resources.zero_trust.finding_event import ZeroTrustFindingEvent from monkey_island.cc.resources.zero_trust.zero_trust_report import ZeroTrustReport diff --git a/monkey/monkey_island/cc/resources/__init__.py b/monkey/monkey_island/cc/resources/__init__.py index bef3988e8..937766e2b 100644 --- a/monkey/monkey_island/cc/resources/__init__.py +++ b/monkey/monkey_island/cc/resources/__init__.py @@ -10,5 +10,4 @@ from .pba_file_upload import PBAFileUpload, LINUX_PBA_TYPE, WINDOWS_PBA_TYPE from .pba_file_download import PBAFileDownload from .agent_events import AgentEvents from .agents import Agents -from .agent_signals import AgentSignals -from .terminate_all_agents import TerminateAllAgents +from .agent_signals import AgentSignals, TerminateAllAgents diff --git a/monkey/monkey_island/cc/resources/agent_signals/__init__.py b/monkey/monkey_island/cc/resources/agent_signals/__init__.py new file mode 100644 index 000000000..b09dbdd87 --- /dev/null +++ b/monkey/monkey_island/cc/resources/agent_signals/__init__.py @@ -0,0 +1,2 @@ +from .agent_signals import AgentSignals +from .terminate_all_agents import TerminateAllAgents diff --git a/monkey/monkey_island/cc/resources/agent_signals.py b/monkey/monkey_island/cc/resources/agent_signals/agent_signals.py similarity index 100% rename from monkey/monkey_island/cc/resources/agent_signals.py rename to monkey/monkey_island/cc/resources/agent_signals/agent_signals.py diff --git a/monkey/monkey_island/cc/resources/terminate_all_agents.py b/monkey/monkey_island/cc/resources/agent_signals/terminate_all_agents.py similarity index 100% rename from monkey/monkey_island/cc/resources/terminate_all_agents.py rename to monkey/monkey_island/cc/resources/agent_signals/terminate_all_agents.py From 275efb2ab1531c34ed7ed98a1dde04ea7eda5f2f Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Thu, 22 Sep 2022 19:23:36 +0000 Subject: [PATCH 42/67] UT: Test on_terminate_agents_signal --- monkey/monkey_island/cc/services/__init__.py | 2 +- .../cc/services/test_agent_signals_service.py | 27 ++++++++++++++++++- 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/services/__init__.py b/monkey/monkey_island/cc/services/__init__.py index c73aff356..8fdb5fc77 100644 --- a/monkey/monkey_island/cc/services/__init__.py +++ b/monkey/monkey_island/cc/services/__init__.py @@ -1,4 +1,4 @@ +from .agent_signals_service import AgentSignalsService from .authentication_service import AuthenticationService from .aws import AWSService -from .agent_signals_service import AgentSignalsService diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_agent_signals_service.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_agent_signals_service.py index 16aab943f..1beb53940 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/test_agent_signals_service.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/test_agent_signals_service.py @@ -4,7 +4,7 @@ from uuid import UUID import pytest from common.types import AgentID -from monkey_island.cc.models import Agent, Simulation +from monkey_island.cc.models import Agent, IslandMode, Simulation from monkey_island.cc.repository import IAgentRepository, ISimulationRepository, UnknownRecordError from monkey_island.cc.services import AgentSignalsService @@ -117,3 +117,28 @@ def test_progenitor_started_before_terminate( signals = agent_signals_service.get_signals(agent.id) assert signals.terminate.timestamp() == TERMINATE_TIMESTAMP + + +def test_on_terminate_agents_signal__stores_timestamp( + agent_signals_service: AgentSignalsService, mock_simulation_repository: ISimulationRepository +): + timestamp = 100 + mock_simulation_repository.get_simulation = MagicMock(return_value=Simulation()) + agent_signals_service.on_terminate_agents_signal(timestamp) + + expected_value = Simulation(terminate_signal_time=timestamp) + assert mock_simulation_repository.save_simulation.called_once_with(expected_value) + + +def test_on_terminate_agents_signal__updates_timestamp( + agent_signals_service: AgentSignalsService, mock_simulation_repository: ISimulationRepository +): + timestamp = 100 + mock_simulation_repository.get_simulation = MagicMock( + return_value=Simulation(mode=IslandMode.RANSOMWARE, terminate_signal_time=50) + ) + + agent_signals_service.on_terminate_agents_signal(timestamp) + + expected_value = Simulation(mode=IslandMode.RANSOMWARE, terminate_signal_time=timestamp) + assert mock_simulation_repository.save_simulation.called_once_with(expected_value) From 8e45a71a155ee3ed8e37e5029be36b68173c7c57 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 22 Sep 2022 12:59:51 -0400 Subject: [PATCH 43/67] Island: Change agent parameter to agent_id in get_signals() --- monkey/monkey_island/cc/services/agent_signals_service.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/monkey/monkey_island/cc/services/agent_signals_service.py b/monkey/monkey_island/cc/services/agent_signals_service.py index 47c677271..4a3c4ef74 100644 --- a/monkey/monkey_island/cc/services/agent_signals_service.py +++ b/monkey/monkey_island/cc/services/agent_signals_service.py @@ -1,18 +1,19 @@ from datetime import datetime -from monkey_island.cc.models import Agent, AgentSignals +from common.types import AgentID from monkey_island.cc.repository import ISimulationRepository +from monkey_island.cc.models import AgentSignals class AgentSignalsService: def __init__(self, simulation_repository: ISimulationRepository): self._simulation_repository = simulation_repository - def get_signals(self, agent: Agent) -> AgentSignals: + def get_signals(self, agent_id: AgentID) -> AgentSignals: """ Gets the signals sent to a particular agent - :param agent: The agent whose signals need to be retrieved + :param agent_id: The ID of the agent whose signals need to be retrieved :return: Signals sent to the relevant agent """ return AgentSignals(timestamp=datetime.now()) From 1dc72e45e7d5fb22067df6e17aabda45d276e302 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 18:45:33 +0530 Subject: [PATCH 44/67] UT: Remove unnecessary wrapper around pytest fixture flask_client in test_terminate_all_agents.py and test_agent_signals.py --- .../cc/resources/test_agent_signals.py | 18 +++++---------- .../cc/resources/test_terminate_all_agents.py | 22 ++++++------------- 2 files changed, 12 insertions(+), 28 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py index 3f644e051..e7a20a9c9 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py @@ -25,20 +25,12 @@ def mock_agent_signals_service(): @pytest.fixture -def flask_client_builder(build_flask_client, mock_agent_signals_service): - def inner(side_effect=None): - container = StubDIContainer() - container.register_instance(AgentSignalsService, mock_agent_signals_service) +def flask_client(build_flask_client, mock_agent_signals_service): + container = StubDIContainer() + container.register_instance(AgentSignalsService, mock_agent_signals_service) - with build_flask_client(container) as flask_client: - return flask_client - - return inner - - -@pytest.fixture -def flask_client(flask_client_builder): - return flask_client_builder() + with build_flask_client(container) as flask_client: + yield flask_client @pytest.mark.parametrize( diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_terminate_all_agents.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_terminate_all_agents.py index 305293f26..bcf91afd4 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_terminate_all_agents.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_terminate_all_agents.py @@ -11,23 +11,15 @@ TIMESTAMP = 123456789 @pytest.fixture -def flask_client_builder(build_flask_client): - def inner(side_effect=None): - container = StubDIContainer() +def flask_client(build_flask_client): + container = StubDIContainer() - mock_island_event_queue = MagicMock(spec=IIslandEventQueue) - mock_island_event_queue.publish.side_effect = side_effect - container.register_instance(IIslandEventQueue, mock_island_event_queue) + mock_island_event_queue = MagicMock(spec=IIslandEventQueue) + mock_island_event_queue.publish.side_effect = None + container.register_instance(IIslandEventQueue, mock_island_event_queue) - with build_flask_client(container) as flask_client: - return flask_client - - return inner - - -@pytest.fixture -def flask_client(flask_client_builder): - return flask_client_builder() + with build_flask_client(container) as flask_client: + yield flask_client def test_terminate_all_agents_post(flask_client): From f9306cf8f157802998986a0e9ac19fd6d6396494 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 09:41:31 +0200 Subject: [PATCH 45/67] Island: Keep naming consistency in island_event_handlers --- monkey/monkey_island/cc/setup/island_event_handlers.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/setup/island_event_handlers.py b/monkey/monkey_island/cc/setup/island_event_handlers.py index 31e7908f6..96d2e1941 100644 --- a/monkey/monkey_island/cc/setup/island_event_handlers.py +++ b/monkey/monkey_island/cc/setup/island_event_handlers.py @@ -26,7 +26,7 @@ def setup_island_event_handlers(container: DIContainer): _subscribe_reset_agent_configuration_events(island_event_queue, container) _subscribe_clear_simulation_data_events(island_event_queue, container) _subscribe_set_island_mode_events(island_event_queue, container) - _subscribe_on_terminate_agents_signal(island_event_queue, container) + _subscribe_terminate_agents_events(island_event_queue, container) def _subscribe_agent_registration_events( @@ -78,7 +78,7 @@ def _subscribe_set_island_mode_events( island_event_queue.subscribe(topic, simulation_repository.set_mode) -def _subscribe_on_terminate_agents_signal( +def _subscribe_terminate_agents_events( island_event_queue: IIslandEventQueue, container: DIContainer ): topic = IslandEventTopic.TERMINATE_AGENTS From a04a6a3cea0c5f6b06f5626eae3fc4a33c74662f Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 22 Sep 2022 13:35:16 -0400 Subject: [PATCH 46/67] Island: Implement AgentSignalsService.get_signals() --- .../cc/services/agent_signals_service.py | 27 +++- .../cc/services/test_agent_signals_service.py | 119 ++++++++++++++++++ 2 files changed, 143 insertions(+), 3 deletions(-) create mode 100644 monkey/tests/unit_tests/monkey_island/cc/services/test_agent_signals_service.py diff --git a/monkey/monkey_island/cc/services/agent_signals_service.py b/monkey/monkey_island/cc/services/agent_signals_service.py index 4a3c4ef74..23a15dba9 100644 --- a/monkey/monkey_island/cc/services/agent_signals_service.py +++ b/monkey/monkey_island/cc/services/agent_signals_service.py @@ -1,13 +1,17 @@ from datetime import datetime +from typing import Optional from common.types import AgentID -from monkey_island.cc.repository import ISimulationRepository from monkey_island.cc.models import AgentSignals +from monkey_island.cc.repository import IAgentRepository, ISimulationRepository class AgentSignalsService: - def __init__(self, simulation_repository: ISimulationRepository): + def __init__( + self, simulation_repository: ISimulationRepository, agent_repository: IAgentRepository + ): self._simulation_repository = simulation_repository + self._agent_repository = agent_repository def get_signals(self, agent_id: AgentID) -> AgentSignals: """ @@ -16,7 +20,24 @@ class AgentSignalsService: :param agent_id: The ID of the agent whose signals need to be retrieved :return: Signals sent to the relevant agent """ - return AgentSignals(timestamp=datetime.now()) + terminate_timestamp = self._get_terminate_signal_timestamp(agent_id) + return AgentSignals(terminate=terminate_timestamp) + + def _get_terminate_signal_timestamp(self, agent_id: AgentID) -> Optional[datetime]: + simulation = self._simulation_repository.get_simulation() + terminate_all_signal_time = simulation.terminate_signal_time + if terminate_all_signal_time is None: + return None + + agent = self._agent_repository.get_agent_by_id(agent_id) + if agent.start_time <= terminate_all_signal_time: + return terminate_all_signal_time + + progenitor = self._agent_repository.get_progenitor(agent) + if progenitor.start_time <= terminate_all_signal_time: + return terminate_all_signal_time + + return None def on_terminate_agents_signal(self, timestamp: datetime): """ diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_agent_signals_service.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_agent_signals_service.py new file mode 100644 index 000000000..16aab943f --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/services/test_agent_signals_service.py @@ -0,0 +1,119 @@ +from unittest.mock import MagicMock +from uuid import UUID + +import pytest + +from common.types import AgentID +from monkey_island.cc.models import Agent, Simulation +from monkey_island.cc.repository import IAgentRepository, ISimulationRepository, UnknownRecordError +from monkey_island.cc.services import AgentSignalsService + +AGENT_1 = Agent( + id=UUID("f811ad00-5a68-4437-bd51-7b5cc1768ad5"), + machine_id=1, + start_time=100, + parent_id=None, +) + +AGENT_2 = Agent( + id=UUID("012e7238-7b81-4108-8c7f-0787bc3f3c10"), + machine_id=2, + start_time=200, + parent_id=AGENT_1.id, +) + +AGENT_3 = Agent( + id=UUID("0fc9afcb-1902-436b-bd5c-1ad194252484"), + machine_id=3, + start_time=300, + parent_id=AGENT_2.id, +) +AGENTS = [AGENT_1, AGENT_2, AGENT_3] + + +@pytest.fixture +def mock_simulation_repository() -> IAgentRepository: + return MagicMock(spec=ISimulationRepository) + + +@pytest.fixture(scope="session") +def mock_agent_repository() -> IAgentRepository: + def get_agent_by_id(agent_id: AgentID) -> Agent: + for agent in AGENTS: + if agent.id == agent_id: + return agent + + raise UnknownRecordError(str(agent_id)) + + agent_repository = MagicMock(spec=IAgentRepository) + agent_repository.get_progenitor = MagicMock(return_value=AGENT_1) + agent_repository.get_agent_by_id = MagicMock(side_effect=get_agent_by_id) + + return agent_repository + + +@pytest.fixture +def agent_signals_service(mock_simulation_repository, mock_agent_repository) -> AgentSignalsService: + return AgentSignalsService(mock_simulation_repository, mock_agent_repository) + + +@pytest.mark.parametrize("agent", AGENTS) +def test_terminate_is_none( + agent, + agent_signals_service: AgentSignalsService, + mock_simulation_repository: ISimulationRepository, +): + mock_simulation_repository.get_simulation = MagicMock( + return_value=Simulation(terminate_signal_time=None) + ) + + signals = agent_signals_service.get_signals(agent.id) + assert signals.terminate is None + + +@pytest.mark.parametrize("agent", AGENTS) +def test_agent_started_before_terminate( + agent, + agent_signals_service: AgentSignalsService, + mock_simulation_repository: ISimulationRepository, +): + TERMINATE_TIMESTAMP = 400 + mock_simulation_repository.get_simulation = MagicMock( + return_value=Simulation(terminate_signal_time=TERMINATE_TIMESTAMP) + ) + + signals = agent_signals_service.get_signals(agent.id) + + assert signals.terminate.timestamp() == TERMINATE_TIMESTAMP + + +@pytest.mark.parametrize("agent", AGENTS) +def test_agent_started_after_terminate( + agent, + agent_signals_service: AgentSignalsService, + mock_simulation_repository: ISimulationRepository, +): + TERMINATE_TIMESTAMP = 50 + mock_simulation_repository.get_simulation = MagicMock( + return_value=Simulation(terminate_signal_time=TERMINATE_TIMESTAMP) + ) + + signals = agent_signals_service.get_signals(agent.id) + + assert signals.terminate is None + + +@pytest.mark.parametrize("agent", AGENTS) +def test_progenitor_started_before_terminate( + agent, + agent_signals_service: AgentSignalsService, + mock_simulation_repository: ISimulationRepository, +): + TERMINATE_TIMESTAMP = 150 + mock_simulation_repository.get_simulation = MagicMock( + return_value=Simulation(terminate_signal_time=TERMINATE_TIMESTAMP) + ) + + signals = agent_signals_service.get_signals(agent.id) + + assert signals.terminate.timestamp() == TERMINATE_TIMESTAMP From dee28841440110f7c0877594245f97d049d70b30 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 18:47:25 +0530 Subject: [PATCH 47/67] UT: Move test_agent_signals.py and test_terminate_all_agents.py under cc/resources/agent_signals/ --- .../cc/resources/{ => agent_signals}/test_agent_signals.py | 0 .../cc/resources/{ => agent_signals}/test_terminate_all_agents.py | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename monkey/tests/unit_tests/monkey_island/cc/resources/{ => agent_signals}/test_agent_signals.py (100%) rename monkey/tests/unit_tests/monkey_island/cc/resources/{ => agent_signals}/test_terminate_all_agents.py (100%) diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py b/monkey/tests/unit_tests/monkey_island/cc/resources/agent_signals/test_agent_signals.py similarity index 100% rename from monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py rename to monkey/tests/unit_tests/monkey_island/cc/resources/agent_signals/test_agent_signals.py diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_terminate_all_agents.py b/monkey/tests/unit_tests/monkey_island/cc/resources/agent_signals/test_terminate_all_agents.py similarity index 100% rename from monkey/tests/unit_tests/monkey_island/cc/resources/test_terminate_all_agents.py rename to monkey/tests/unit_tests/monkey_island/cc/resources/agent_signals/test_terminate_all_agents.py From a3d94d7a493578cacb23bd21cc657af30a83a5aa Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 13:51:21 +0200 Subject: [PATCH 48/67] Agent: Add get_agent_signals to IIslandAPIClient --- .../island_api_client/i_island_api_client.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/monkey/infection_monkey/island_api_client/i_island_api_client.py b/monkey/infection_monkey/island_api_client/i_island_api_client.py index 8ecd98b49..025261e06 100644 --- a/monkey/infection_monkey/island_api_client/i_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/i_island_api_client.py @@ -1,5 +1,6 @@ from abc import ABC, abstractmethod -from typing import Sequence +from datetime import datetime +from typing import Optional, Sequence from common import AgentRegistrationData, OperatingSystem from common.agent_configuration import AgentConfiguration @@ -143,3 +144,15 @@ class IIslandAPIClient(ABC): :raises IslandAPITimeoutError: If the command timed out :return: Credentials """ + + @abstractmethod + def get_agent_signals(self, agent_id: str) -> Optional[datetime]: + """ + Get agent signals from the island + + :raises IslandAPIConnectionError: If the client could not connect to the island + :raises IslandAPIRequestError: If there was a problem with the client request + :raises IslandAPIRequestFailedError: If the server experienced an error + :raises IslandAPITimeoutError: If the command timed out + :return: Terminate datetime + """ From 296f4e55dff04fa84402b303a3f2dc34c3c16f8f Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 17:07:15 +0200 Subject: [PATCH 49/67] Common, Island: Move AgentSignals model to Common --- monkey/common/__init__.py | 1 + monkey/{monkey_island/cc/models => common}/agent_signals.py | 2 +- monkey/monkey_island/cc/models/__init__.py | 1 - 3 files changed, 2 insertions(+), 2 deletions(-) rename monkey/{monkey_island/cc/models => common}/agent_signals.py (71%) diff --git a/monkey/common/__init__.py b/monkey/common/__init__.py index 4898b2df8..69190bce5 100644 --- a/monkey/common/__init__.py +++ b/monkey/common/__init__.py @@ -7,3 +7,4 @@ from .operating_system import OperatingSystem from . import types from . import base_models from .agent_registration_data import AgentRegistrationData +from .agent_signals import AgentSignals diff --git a/monkey/monkey_island/cc/models/agent_signals.py b/monkey/common/agent_signals.py similarity index 71% rename from monkey/monkey_island/cc/models/agent_signals.py rename to monkey/common/agent_signals.py index 37af7b4c1..c351e8823 100644 --- a/monkey/monkey_island/cc/models/agent_signals.py +++ b/monkey/common/agent_signals.py @@ -1,7 +1,7 @@ from datetime import datetime from typing import Optional -from common.base_models import InfectionMonkeyBaseModel +from .base_models import InfectionMonkeyBaseModel class AgentSignals(InfectionMonkeyBaseModel): diff --git a/monkey/monkey_island/cc/models/__init__.py b/monkey/monkey_island/cc/models/__init__.py index ca4078faa..65e63fe14 100644 --- a/monkey/monkey_island/cc/models/__init__.py +++ b/monkey/monkey_island/cc/models/__init__.py @@ -15,4 +15,3 @@ from .communication_type import CommunicationType from .node import Node from common.types import AgentID from .agent import Agent -from .agent_signals import AgentSignals From 88c011e883a0336f3beebb6686a2dca961937138 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 13:51:50 +0200 Subject: [PATCH 50/67] Agent: Implement IIslandAPIClient.get_agent_signals in HTTPIslandAPIClient --- .../http_island_api_client.py | 15 ++++- .../test_http_island_api_client.py | 60 +++++++++++++++++++ 2 files changed, 74 insertions(+), 1 deletion(-) diff --git a/monkey/infection_monkey/island_api_client/http_island_api_client.py b/monkey/infection_monkey/island_api_client/http_island_api_client.py index 6cd1d86a1..5c65ebf6f 100644 --- a/monkey/infection_monkey/island_api_client/http_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/http_island_api_client.py @@ -1,8 +1,9 @@ import functools import json import logging +from datetime import datetime from pprint import pformat -from typing import List, Sequence +from typing import List, Optional, Sequence import requests @@ -199,6 +200,18 @@ class HTTPIslandAPIClient(IIslandAPIClient): return serialized_events + @handle_island_errors + @convert_json_error_to_island_api_error + def get_agent_signals(self, agent_id: str) -> Optional[datetime]: + url = f"{self._api_url}/agent-signals/{agent_id}" + response = requests.get( # noqa: DUO123 + url, + verify=False, + timeout=SHORT_REQUEST_TIMEOUT, + ) + response.raise_for_status() + return response.json()["terminate"] + class HTTPIslandAPIClientFactory(AbstractIslandAPIClientFactory): def __init__( diff --git a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py index 03117b006..94ad23cac 100644 --- a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py +++ b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py @@ -33,6 +33,8 @@ AGENT_REGISTRATION = AgentRegistrationData( network_interfaces=[], ) +TIMESTAMP = 123456789 + ISLAND_URI = f"https://{SERVER}/api?action=is-up" ISLAND_SEND_LOG_URI = f"https://{SERVER}/api/log" ISLAND_GET_PBA_FILE_URI = f"https://{SERVER}/api/pba/download/{PBA_FILE}" @@ -42,6 +44,7 @@ ISLAND_REGISTER_AGENT_URI = f"https://{SERVER}/api/agents" ISLAND_AGENT_STOP_URI = f"https://{SERVER}/api/monkey-control/needs-to-stop/{AGENT_ID}" ISLAND_GET_CONFIG_URI = f"https://{SERVER}/api/agent-configuration" ISLAND_GET_PROPAGATION_CREDENTIALS_URI = f"https://{SERVER}/api/propagation-credentials" +ISLAND_GET_AGENT_SIGNALS = f"https://{SERVER}/api/agent-signals/{AGENT_ID}" class Event1(AbstractAgentEvent): @@ -461,3 +464,60 @@ def test_island_api_client_get_credentials_for_propagation__bad_json(island_api_ with pytest.raises(IslandAPIRequestFailedError): m.get(ISLAND_GET_PROPAGATION_CREDENTIALS_URI, content=b"bad") island_api_client.get_credentials_for_propagation() + + +@pytest.mark.parametrize( + "actual_error, expected_error", + [ + (requests.exceptions.ConnectionError, IslandAPIConnectionError), + (TimeoutError, IslandAPITimeoutError), + ], +) +def test_island_api_client__get_agent_signals(island_api_client, actual_error, expected_error): + with requests_mock.Mocker() as m: + m.get(ISLAND_URI) + island_api_client.connect(SERVER) + + with pytest.raises(expected_error): + m.get(ISLAND_GET_AGENT_SIGNALS, exc=actual_error) + island_api_client.get_agent_signals(agent_id=AGENT_ID) + + +@pytest.mark.parametrize( + "status_code, expected_error", + [ + (401, IslandAPIRequestError), + (501, IslandAPIRequestFailedError), + ], +) +def test_island_api_client_get_agent_signals__status_code( + island_api_client, status_code, expected_error +): + with requests_mock.Mocker() as m: + m.get(ISLAND_URI) + island_api_client.connect(SERVER) + + with pytest.raises(expected_error): + m.get(ISLAND_GET_AGENT_SIGNALS, status_code=status_code) + island_api_client.get_agent_signals(agent_id=AGENT_ID) + + +def test_island_api_client_get_agent_signals(island_api_client): + with requests_mock.Mocker() as m: + m.get(ISLAND_URI) + island_api_client.connect(SERVER) + + m.get(ISLAND_GET_AGENT_SIGNALS, json={"terminate": TIMESTAMP}) + actual_terminate_timestamp = island_api_client.get_agent_signals(agent_id=AGENT_ID) + + assert actual_terminate_timestamp == TIMESTAMP + + +def test_island_api_client_get_agent_signals__bad_json(island_api_client): + with requests_mock.Mocker() as m: + m.get(ISLAND_URI) + island_api_client.connect(SERVER) + + with pytest.raises(IslandAPIError): + m.get(ISLAND_GET_AGENT_SIGNALS, json={"bogus": "vogus"}) + island_api_client.get_agent_signals(agent_id=AGENT_ID) From ffa5f90cbd6341260fc5f73778ebc020126f9f00 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 17:08:34 +0200 Subject: [PATCH 51/67] Island: Use common.agent_signals in AgentSignalsService --- monkey/monkey_island/cc/services/agent_signals_service.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/services/agent_signals_service.py b/monkey/monkey_island/cc/services/agent_signals_service.py index 473a10066..bf0b3e61f 100644 --- a/monkey/monkey_island/cc/services/agent_signals_service.py +++ b/monkey/monkey_island/cc/services/agent_signals_service.py @@ -2,8 +2,9 @@ import logging from datetime import datetime from typing import Optional +from common.agent_signals import AgentSignals from common.types import AgentID -from monkey_island.cc.models import AgentSignals, Simulation +from monkey_island.cc.models import Simulation from monkey_island.cc.repository import IAgentRepository, ISimulationRepository logger = logging.getLogger(__name__) From 3da90223fcabee9e288db5d3f83aa42a5a7cd557 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 13:52:26 +0200 Subject: [PATCH 52/67] Agent: Use IIslandAPIClient.get_agent_signals in ControlChannel --- monkey/infection_monkey/master/control_channel.py | 2 +- .../infection_monkey/master/test_control_channel.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index 48b827f63..947d6c0da 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/monkey/infection_monkey/master/control_channel.py @@ -36,7 +36,7 @@ class ControlChannel(IControlChannel): if not self._control_channel_server: logger.error("Agent should stop because it can't connect to the C&C server.") return True - return self._island_api_client.should_agent_stop(self._agent_id) + return self._island_api_client.get_agent_signals(self._agent_id) is not None @handle_island_api_errors def get_config(self) -> AgentConfiguration: diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py b/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py index 658635615..1da0d0713 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py +++ b/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py @@ -35,14 +35,14 @@ def control_channel(island_api_client) -> ControlChannel: def test_control_channel__should_agent_stop(control_channel, island_api_client): control_channel.should_agent_stop() - assert island_api_client.should_agent_stop.called_once() + assert island_api_client.get_agent_signals.called_once() @pytest.mark.parametrize("api_error", CONTROL_CHANNEL_API_ERRORS) def test_control_channel__should_agent_stop_raises_error( control_channel, island_api_client, api_error ): - island_api_client.should_agent_stop.side_effect = api_error() + island_api_client.get_agent_signals.side_effect = api_error() with pytest.raises(IslandCommunicationError): control_channel.should_agent_stop() From 7a9ac1a6bacc7e1a9afc90ead968be30fea99698 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 17:09:11 +0200 Subject: [PATCH 53/67] UT: Fix AgentSignals endpoint tests to use common.agent_signals --- .../cc/resources/agent_signals/test_agent_signals.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/agent_signals/test_agent_signals.py b/monkey/tests/unit_tests/monkey_island/cc/resources/agent_signals/test_agent_signals.py index e7a20a9c9..723ac1afe 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/agent_signals/test_agent_signals.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/agent_signals/test_agent_signals.py @@ -5,7 +5,7 @@ from uuid import UUID import pytest from tests.common import StubDIContainer -from monkey_island.cc.models import AgentSignals as Signals +from common.agent_signals import AgentSignals as Signals from monkey_island.cc.repository import RetrievalError, StorageError from monkey_island.cc.services import AgentSignalsService From 67956358bd802ea65ddb901223a92446948f60b5 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 14:11:57 +0200 Subject: [PATCH 54/67] Agent: Remove shoudl_agent_stop from IIslandAPIClient --- .../island_api_client/i_island_api_client.py | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/monkey/infection_monkey/island_api_client/i_island_api_client.py b/monkey/infection_monkey/island_api_client/i_island_api_client.py index 025261e06..34348aeef 100644 --- a/monkey/infection_monkey/island_api_client/i_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/i_island_api_client.py @@ -108,19 +108,6 @@ class IIslandAPIClient(ABC): :raises IslandAPITimeoutError: If the command timed out """ - @abstractmethod - def should_agent_stop(self, agent_id: str) -> bool: - """ - Check with the island to see if the agent should stop - - :param agent_id: The agent identifier for the agent to check - :raises IslandAPIConnectionError: If the client could not connect to the island - :raises IslandAPIRequestError: If there was a problem with the client request - :raises IslandAPIRequestFailedError: If the server experienced an error - :raises IslandAPITimeoutError: If the command timed out - :return: True if the agent should stop, otherwise False - """ - @abstractmethod def get_config(self) -> AgentConfiguration: """ From a2be330d16841be96ac1f5a92980b94aa93f4346 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 17:10:24 +0200 Subject: [PATCH 55/67] Island: IIslandAPIClient.get_agent_signals to return AgentSignals --- .../island_api_client/i_island_api_client.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/monkey/infection_monkey/island_api_client/i_island_api_client.py b/monkey/infection_monkey/island_api_client/i_island_api_client.py index c2a4dc899..3101cc8f8 100644 --- a/monkey/infection_monkey/island_api_client/i_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/i_island_api_client.py @@ -1,8 +1,7 @@ from abc import ABC, abstractmethod -from datetime import datetime from typing import Optional, Sequence -from common import AgentRegistrationData, OperatingSystem +from common import AgentRegistrationData, AgentSignals, OperatingSystem from common.agent_configuration import AgentConfiguration from common.agent_events import AbstractAgentEvent from common.credentials import Credentials @@ -133,7 +132,7 @@ class IIslandAPIClient(ABC): """ @abstractmethod - def get_agent_signals(self, agent_id: str) -> Optional[datetime]: + def get_agent_signals(self, agent_id: str) -> AgentSignals: """ Gets an agent's signals from the island @@ -142,5 +141,5 @@ class IIslandAPIClient(ABC): :raises IslandAPIRequestError: If there was a problem with the client request :raises IslandAPIRequestFailedError: If the server experienced an error :raises IslandAPITimeoutError: If the command timed out - :return: The relevant agent's terminate signal's timestamp + :return: The relevant agent's signal's """ From edf0593d4a2b29f1d3122849dff545fa3b07ff0b Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 14:12:22 +0200 Subject: [PATCH 56/67] Agent: Remove should_agent_stop from HTTPIslandAPIClient --- .../http_island_api_client.py | 13 ------ .../test_http_island_api_client.py | 46 ------------------- 2 files changed, 59 deletions(-) diff --git a/monkey/infection_monkey/island_api_client/http_island_api_client.py b/monkey/infection_monkey/island_api_client/http_island_api_client.py index 5c65ebf6f..65e34df3c 100644 --- a/monkey/infection_monkey/island_api_client/http_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/http_island_api_client.py @@ -147,19 +147,6 @@ class HTTPIslandAPIClient(IIslandAPIClient): ) response.raise_for_status() - @handle_island_errors - @convert_json_error_to_island_api_error - def should_agent_stop(self, agent_id: str) -> bool: - url = f"{self._api_url}/monkey-control/needs-to-stop/{agent_id}" - response = requests.get( # noqa: DUO123 - url, - verify=False, - timeout=SHORT_REQUEST_TIMEOUT, - ) - response.raise_for_status() - - return response.json()["stop_agent"] - @handle_island_errors @convert_json_error_to_island_api_error def get_config(self) -> AgentConfiguration: diff --git a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py index 94ad23cac..376425696 100644 --- a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py +++ b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py @@ -328,52 +328,6 @@ def test_island_api_client_register_agent__status_code( island_api_client.register_agent(AGENT_REGISTRATION) -@pytest.mark.parametrize( - "actual_error, expected_error", - [ - (requests.exceptions.ConnectionError, IslandAPIConnectionError), - (TimeoutError, IslandAPITimeoutError), - ], -) -def test_island_api_client__should_agent_stop(island_api_client, actual_error, expected_error): - with requests_mock.Mocker() as m: - m.get(ISLAND_URI) - island_api_client.connect(SERVER) - - with pytest.raises(expected_error): - m.get(ISLAND_AGENT_STOP_URI, exc=actual_error) - island_api_client.should_agent_stop(AGENT_ID) - - -@pytest.mark.parametrize( - "status_code, expected_error", - [ - (401, IslandAPIRequestError), - (501, IslandAPIRequestFailedError), - ], -) -def test_island_api_client_should_agent_stop__status_code( - island_api_client, status_code, expected_error -): - with requests_mock.Mocker() as m: - m.get(ISLAND_URI) - island_api_client.connect(SERVER) - - with pytest.raises(expected_error): - m.get(ISLAND_AGENT_STOP_URI, status_code=status_code) - island_api_client.should_agent_stop(AGENT_ID) - - -def test_island_api_client_should_agent_stop__bad_json(island_api_client): - with requests_mock.Mocker() as m: - m.get(ISLAND_URI) - island_api_client.connect(SERVER) - - with pytest.raises(IslandAPIRequestFailedError): - m.get(ISLAND_AGENT_STOP_URI, content=b"bad") - island_api_client.should_agent_stop(AGENT_ID) - - @pytest.mark.parametrize( "actual_error, expected_error", [ From 6299529f4a5dbad50afd9645cb483daa2830b8bb Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 17:11:44 +0200 Subject: [PATCH 57/67] Island: Modify HTTPIslandAPIClient.get_agent_signals to return AgentSignals --- .../island_api_client/http_island_api_client.py | 9 ++++----- .../test_http_island_api_client.py | 13 +++++++------ 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/monkey/infection_monkey/island_api_client/http_island_api_client.py b/monkey/infection_monkey/island_api_client/http_island_api_client.py index 65e34df3c..09fcbf762 100644 --- a/monkey/infection_monkey/island_api_client/http_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/http_island_api_client.py @@ -1,13 +1,12 @@ import functools import json import logging -from datetime import datetime from pprint import pformat -from typing import List, Optional, Sequence +from typing import List, Sequence import requests -from common import AgentRegistrationData, OperatingSystem +from common import AgentRegistrationData, AgentSignals, OperatingSystem from common.agent_configuration import AgentConfiguration from common.agent_event_serializers import AgentEventSerializerRegistry, JSONSerializable from common.agent_events import AbstractAgentEvent @@ -189,7 +188,7 @@ class HTTPIslandAPIClient(IIslandAPIClient): @handle_island_errors @convert_json_error_to_island_api_error - def get_agent_signals(self, agent_id: str) -> Optional[datetime]: + def get_agent_signals(self, agent_id: str) -> AgentSignals: url = f"{self._api_url}/agent-signals/{agent_id}" response = requests.get( # noqa: DUO123 url, @@ -197,7 +196,7 @@ class HTTPIslandAPIClient(IIslandAPIClient): timeout=SHORT_REQUEST_TIMEOUT, ) response.raise_for_status() - return response.json()["terminate"] + return AgentSignals(**response.json()) class HTTPIslandAPIClientFactory(AbstractIslandAPIClientFactory): diff --git a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py index 9505e6649..a861a4849 100644 --- a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py +++ b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py @@ -4,7 +4,7 @@ import pytest import requests import requests_mock -from common import OperatingSystem +from common import AgentSignals, OperatingSystem from common.agent_event_serializers import ( AgentEventSerializerRegistry, PydanticAgentEventSerializer, @@ -456,16 +456,17 @@ def test_island_api_client_get_agent_signals__status_code( island_api_client.get_agent_signals(agent_id=AGENT_ID) -@pytest.mark.parametrize("expected_timestamp", [TIMESTAMP, None]) -def test_island_api_client_get_agent_signals(island_api_client, expected_timestamp): +@pytest.mark.parametrize("timestamp", [TIMESTAMP, None]) +def test_island_api_client_get_agent_signals(island_api_client, timestamp): + expected_agent_signals = AgentSignals(terminate=timestamp) with requests_mock.Mocker() as m: m.get(ISLAND_URI) island_api_client.connect(SERVER) - m.get(ISLAND_GET_AGENT_SIGNALS, json={"terminate": expected_timestamp}) - actual_terminate_timestamp = island_api_client.get_agent_signals(agent_id=AGENT_ID) + m.get(ISLAND_GET_AGENT_SIGNALS, json={"terminate": timestamp}) + actual_agent_signals = island_api_client.get_agent_signals(agent_id=AGENT_ID) - assert actual_terminate_timestamp == expected_timestamp + assert actual_agent_signals == expected_agent_signals def test_island_api_client_get_agent_signals__bad_json(island_api_client): From d1fc4fa7f42f7e86e5f8fd6b5077b6b834b56477 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 15:49:41 +0200 Subject: [PATCH 58/67] UT: Parametrize HTTPIslandAPIClient get_agent_signals test --- .../island_api_client/test_http_island_api_client.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py index 376425696..9505e6649 100644 --- a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py +++ b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py @@ -456,15 +456,16 @@ def test_island_api_client_get_agent_signals__status_code( island_api_client.get_agent_signals(agent_id=AGENT_ID) -def test_island_api_client_get_agent_signals(island_api_client): +@pytest.mark.parametrize("expected_timestamp", [TIMESTAMP, None]) +def test_island_api_client_get_agent_signals(island_api_client, expected_timestamp): with requests_mock.Mocker() as m: m.get(ISLAND_URI) island_api_client.connect(SERVER) - m.get(ISLAND_GET_AGENT_SIGNALS, json={"terminate": TIMESTAMP}) + m.get(ISLAND_GET_AGENT_SIGNALS, json={"terminate": expected_timestamp}) actual_terminate_timestamp = island_api_client.get_agent_signals(agent_id=AGENT_ID) - assert actual_terminate_timestamp == TIMESTAMP + assert actual_terminate_timestamp == expected_timestamp def test_island_api_client_get_agent_signals__bad_json(island_api_client): From e5c5cce94e530b1a1671b577dace5363452520c4 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 17:12:39 +0200 Subject: [PATCH 59/67] Agent: Modify should_agent_stop to use AgentSignals model --- monkey/infection_monkey/master/control_channel.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index 947d6c0da..ad2081065 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/monkey/infection_monkey/master/control_channel.py @@ -36,7 +36,8 @@ class ControlChannel(IControlChannel): if not self._control_channel_server: logger.error("Agent should stop because it can't connect to the C&C server.") return True - return self._island_api_client.get_agent_signals(self._agent_id) is not None + agent_signals = self._island_api_client.get_agent_signals(self._agent_id) + return agent_signals.terminate is not None @handle_island_api_errors def get_config(self) -> AgentConfiguration: From a314efb8d963d5c617fb844ead64d138ba1f33e4 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 15:55:12 +0200 Subject: [PATCH 60/67] Agent: Reword get_agent_signals docstring --- .../island_api_client/i_island_api_client.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/island_api_client/i_island_api_client.py b/monkey/infection_monkey/island_api_client/i_island_api_client.py index 34348aeef..c2a4dc899 100644 --- a/monkey/infection_monkey/island_api_client/i_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/i_island_api_client.py @@ -135,11 +135,12 @@ class IIslandAPIClient(ABC): @abstractmethod def get_agent_signals(self, agent_id: str) -> Optional[datetime]: """ - Get agent signals from the island + Gets an agent's signals from the island + :param agent_id: ID of the agent whose signals should be retrieved :raises IslandAPIConnectionError: If the client could not connect to the island :raises IslandAPIRequestError: If there was a problem with the client request :raises IslandAPIRequestFailedError: If the server experienced an error :raises IslandAPITimeoutError: If the command timed out - :return: Terminate datetime + :return: The relevant agent's terminate signal's timestamp """ From a5f1117ce37b705338c122415b4ab7d4863436d6 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 23 Sep 2022 12:08:56 -0400 Subject: [PATCH 61/67] Island: Fix grammar in docstring Co-authored-by: Kekoa Kaaikala --- .../infection_monkey/island_api_client/i_island_api_client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/infection_monkey/island_api_client/i_island_api_client.py b/monkey/infection_monkey/island_api_client/i_island_api_client.py index 3101cc8f8..2cf1e72ab 100644 --- a/monkey/infection_monkey/island_api_client/i_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/i_island_api_client.py @@ -141,5 +141,5 @@ class IIslandAPIClient(ABC): :raises IslandAPIRequestError: If there was a problem with the client request :raises IslandAPIRequestFailedError: If the server experienced an error :raises IslandAPITimeoutError: If the command timed out - :return: The relevant agent's signal's + :return: The relevant agent's signals """ From f7198ea98ace6b122435dc4e5db33db9230be68a Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 23 Sep 2022 12:25:09 -0400 Subject: [PATCH 62/67] UT: Add proper test for ControlChannel.should_agent_stop() --- .../master/test_control_channel.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py b/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py index 1da0d0713..efc52f79f 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py +++ b/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py @@ -1,8 +1,10 @@ +from typing import Optional from unittest.mock import MagicMock import pytest -from infection_monkey.i_control_channel import IslandCommunicationError +from common import AgentSignals +from infection_monkey.i_control_channel import IControlChannel, IslandCommunicationError from infection_monkey.island_api_client import ( IIslandAPIClient, IslandAPIConnectionError, @@ -33,9 +35,17 @@ def control_channel(island_api_client) -> ControlChannel: return ControlChannel(SERVER, AGENT_ID, island_api_client) -def test_control_channel__should_agent_stop(control_channel, island_api_client): - control_channel.should_agent_stop() - assert island_api_client.get_agent_signals.called_once() +@pytest.mark.parametrize("signal_time,expected_should_stop", [(1663950115, True), (None, False)]) +def test_control_channel__should_agent_stop( + control_channel: IControlChannel, + island_api_client: IIslandAPIClient, + signal_time: Optional[int], + expected_should_stop: bool, +): + island_api_client.get_agent_signals = MagicMock( + return_value=AgentSignals(terminate=signal_time) + ) + assert control_channel.should_agent_stop() is expected_should_stop @pytest.mark.parametrize("api_error", CONTROL_CHANNEL_API_ERRORS) From feb8288c98629b0f4a10f54c006843e11ebd4dfb Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 23 Sep 2022 11:25:08 -0400 Subject: [PATCH 63/67] Agent: Pass the correct agent ID to ControlChannel --- monkey/infection_monkey/monkey.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index b5c20fd76..3db3cb856 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -123,7 +123,7 @@ class InfectionMonkey: self._control_client = ControlClient( server_address=server, island_api_client=self._island_api_client ) - self._control_channel = ControlChannel(server, GUID, self._island_api_client) + self._control_channel = ControlChannel(server, get_agent_id(), self._island_api_client) self._register_agent(server) # TODO Refactor the telemetry messengers to accept control client From b11cd9c5f17843cc8c94e40dc4f4386f78a9e11b Mon Sep 17 00:00:00 2001 From: vakarisz Date: Fri, 23 Sep 2022 12:08:24 +0300 Subject: [PATCH 64/67] 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) From 3fbbc01861f8ec5bbc1fcd637995aff071f523b6 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 23 Sep 2022 12:47:11 -0400 Subject: [PATCH 65/67] Changelog: Add entries for #2261 --- CHANGELOG.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index a01568e86..72782b3e4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,7 @@ Changelog](https://keepachangelog.com/en/1.0.0/). - The ability to customize the file extension used by ransomware when encrypting files. #1242 - `/api/agents` endpoint. +- `/api/agent-signals` endpoint. #2261 ### Changed - Reset workflow. Now it's possible to delete data gathered by agents without @@ -64,6 +65,7 @@ Changelog](https://keepachangelog.com/en/1.0.0/). - Tunneling to relays to provide better firewall evasion, faster Island connection times, unlimited hops, and a more resilient way for agents to call home. #2216, #1583 +- "/api/monkey-control/stop-all-agents" to "/api/agent-signals/terminate-all-agents" #2261 ### Removed - VSFTPD exploiter. #1533 @@ -109,6 +111,7 @@ Changelog](https://keepachangelog.com/en/1.0.0/). - "/api/configuration/export" endpoint. #2002 - "/api/island-configuration" endpoint. #2003 - "-t/--tunnel" from agent command line arguments. #2216 +- "/api/monkey-control/neets-to-stop". #2261 ### Fixed - A bug in network map page that caused delay of telemetry log loading. #1545 From 6c63d4edbd4c5d83ee0a8e27de573f8ccfc90316 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 23 Sep 2022 12:57:09 -0400 Subject: [PATCH 66/67] Agent: Remove unused "Option" from i_island_api_client.py --- .../infection_monkey/island_api_client/i_island_api_client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/infection_monkey/island_api_client/i_island_api_client.py b/monkey/infection_monkey/island_api_client/i_island_api_client.py index 2cf1e72ab..e29c19115 100644 --- a/monkey/infection_monkey/island_api_client/i_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/i_island_api_client.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from typing import Optional, Sequence +from typing import Sequence from common import AgentRegistrationData, AgentSignals, OperatingSystem from common.agent_configuration import AgentConfiguration From 73841fb04e4095028d060ebaf2f6e3c72f7af834 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 23 Sep 2022 13:04:21 -0400 Subject: [PATCH 67/67] Project: Update vulture_allowlist.py --- vulture_allowlist.py | 24 ++++++++++-------------- 1 file changed, 10 insertions(+), 14 deletions(-) diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 4efe4a6b4..aa5c8127c 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -8,14 +8,15 @@ from common.agent_configuration.agent_sub_configurations import ( ScanTargetConfiguration, ) from common.credentials import Credentials, LMHash, NTHash +from common.types import SocketAddress from infection_monkey.exploit.log4shell_utils.ldap_server import LDAPServerFactory from monkey_island.cc.event_queue import IslandEventTopic, PyPubSubIslandEventQueue -from monkey_island.cc.models import Report, Simulation -from monkey_island.cc.models import AgentSignals, Report +from monkey_island.cc.models import Report from monkey_island.cc.models.networkmap import Arc, NetworkMap from monkey_island.cc.repository import MongoAgentRepository, MongoMachineRepository from monkey_island.cc.repository.attack.IMitigationsRepository import IMitigationsRepository from monkey_island.cc.repository.i_agent_event_repository import IAgentEventRepository +from monkey_island.cc.repository.i_agent_log_repository import IAgentLogRepository from monkey_island.cc.repository.i_agent_repository import IAgentRepository from monkey_island.cc.repository.i_attack_repository import IAttackRepository from monkey_island.cc.repository.i_config_repository import IConfigRepository @@ -246,8 +247,6 @@ IMitigationsRepository.save_mitigations IAgentRepository.upsert_agent IAgentRepository.get_agent_by_id IAgentRepository.get_running_agents -IAgentRepository.get_progenitor -descendant agent IAttackRepository.get_attack_report IAttackRepository.save_attack_report @@ -304,6 +303,11 @@ IAgentEventRepository.get_events_by_type IAgentEventRepository.get_events_by_tag IAgentEventRepository.get_events_by_source +# TODO: Remove once #2274 is closed +IAgentLogRepository +IAgentLogRepository.upsert_agent_log +IAgentLogRepository.get_agent_log + # pydantic base models underscore_attrs_are_private @@ -323,13 +327,5 @@ EXPLOITED CC CC_TUNNEL -IslandEventTopic.AGENT_CONNECTED -IslandEventTopic.CLEAR_SIMULATION_DATA -IslandEventTopic.RESET_AGENT_CONFIGURATION -# TODO: Remove after #2261 is closed -IslandEventTopic.TERMINATE_AGENTS - -Simulation.terminate_signal_time - -AgentSignalsService.get_signals -AgentSignalsService.on_terminate_agents_signal +# TODO: Remove after #2323 +SocketAddress