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 == {}