From 0775449fa95bbabc8d779f5b93778deef0c5ccfa Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Thu, 22 Sep 2022 17:42:43 +0000 Subject: [PATCH 01/33] 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 02/33] 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 03/33] 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 04/33] 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 05/33] 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 06/33] 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 07/33] 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 08/33] 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 09/33] 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 10/33] 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 11/33] 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 12/33] 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 13/33] 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 14/33] 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 15/33] 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 16/33] 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 17/33] 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 18/33] 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 19/33] 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 20/33] 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 21/33] 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 22/33] 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 23/33] 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 24/33] 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 25/33] 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 26/33] 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 27/33] 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 b666078e7d1743b5be9faef012fc553c6a31141a Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 18:31:37 +0530 Subject: [PATCH 28/33] 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 850857c8a1bf3ef972789b522d91a918d0b30206 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 18:32:18 +0530 Subject: [PATCH 29/33] 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 6174e8dfcbdfbf4d97f1dd0149a4dd3dadaf8894 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 18:32:27 +0530 Subject: [PATCH 30/33] 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 c4642141f0358e95a354300ac8414c1d565e038e Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 18:36:03 +0530 Subject: [PATCH 31/33] 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 1dc72e45e7d5fb22067df6e17aabda45d276e302 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 18:45:33 +0530 Subject: [PATCH 32/33] 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 dee28841440110f7c0877594245f97d049d70b30 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Fri, 23 Sep 2022 18:47:25 +0530 Subject: [PATCH 33/33] 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