Island: Add AgentSignals resource

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)

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

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/<string:agent_id>"]
+
+    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

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