Merge pull request #2336 from guardicore/2261-add-agent-signals-resource

2261 add agent signals resource

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-agents", json={"terminate_time": time.time()}
         )
         if response.ok:
             LOGGER.info("Killing all monkeys after the test.")

monkey/monkey_island/cc/app.py

@@ -15,6 +15,7 @@ from monkey_island.cc.resources import (
     AgentConfiguration,
     AgentEvents,
     Agents,
+    AgentSignals,
     ClearSimulationData,
     IPAddresses,
     IslandLog,
@@ -23,9 +24,10 @@ 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, 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
@@ -188,6 +190,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)
@@ -197,7 +200,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
@@ -211,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):

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, TerminateAllAgents

monkey/monkey_island/cc/resources/agent_controls/__init__.py

@@ -1,2 +1 @@
-from .stop_all_agents import StopAllAgents
 from .stop_agent_check import StopAgentCheck

monkey/monkey_island/cc/resources/agent_controls/stop_all_agents.py

@@ -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)}

monkey/monkey_island/cc/resources/agent_signals/__init__.py

@@ -0,0 +1,2 @@
+from .agent_signals import AgentSignals
+from .terminate_all_agents import TerminateAllAgents

monkey/monkey_island/cc/resources/agent_signals/agent_signals.py

@@ -0,0 +1,21 @@
+import logging
+from http import HTTPStatus
+
+from monkey_island.cc.resources.AbstractResource import AbstractResource
+from monkey_island.cc.services import AgentSignalsService
+
+logger = logging.getLogger(__name__)
+
+
+class AgentSignals(AbstractResource):
+    urls = ["/api/agent-signals/<string:agent_id>"]
+
+    def __init__(
+        self,
+        agent_signals_service: AgentSignalsService,
+    ):
+        self._agent_signals_service = agent_signals_service
+
+    def get(self, agent_id: str):
+        agent_signals = self._agent_signals_service.get_signals(agent_id)
+        return agent_signals.dict(simplify=True), HTTPStatus.OK

monkey/monkey_island/cc/resources/agent_signals/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/agent-signals/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

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)
@@ -37,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():

monkey/monkey_island/cc/ui/src/components/pages/MapPage.js

@@ -84,12 +84,12 @@ class MapPageComponent extends AuthComponent {
   }
 
   killAllMonkeys = () => {
-    this.authFetch('/api/monkey-control/stop-all-agents',
+    this.authFetch('/api/agent-signals/terminate-all-agents',
       {
         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})});

monkey/tests/unit_tests/monkey_island/cc/resources/agent_signals/test_agent_signals.py

@@ -0,0 +1,62 @@
+from http import HTTPStatus
+from unittest.mock import MagicMock
+from uuid import UUID
+
+import pytest
+from tests.common import StubDIContainer
+
+from monkey_island.cc.models import AgentSignals as Signals
+from monkey_island.cc.repository import RetrievalError, StorageError
+from monkey_island.cc.services import AgentSignalsService
+
+TIMESTAMP_1 = 123456789
+TIMESTAMP_2 = 123546789
+
+SIGNALS_1 = Signals(terminate=TIMESTAMP_1)
+SIGNALS_2 = Signals(terminate=TIMESTAMP_2)
+
+AGENT_ID_1 = UUID("c0dd10b3-e21a-4da9-9d96-a99c19ebd7c5")
+AGENT_ID_2 = UUID("9b4279f6-6ec5-4953-821e-893ddc71a988")
+
+
+@pytest.fixture
+def mock_agent_signals_service():
+    return MagicMock(spec=AgentSignalsService)
+
+
+@pytest.fixture
+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:
+        yield flask_client
+
+
+@pytest.mark.parametrize(
+    "url, signals",
+    [
+        (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
+    resp = flask_client.get(url, follow_redirects=True)
+    assert resp.status_code == HTTPStatus.OK
+    assert resp.json == signals.dict(simplify=True)
+
+
+@pytest.mark.parametrize(
+    "url, error",
+    [
+        (f"/api/agent-signals/{AGENT_ID_1}", RetrievalError),
+        (f"/api/agent-signals/{AGENT_ID_2}", 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

monkey/tests/unit_tests/monkey_island/cc/resources/agent_signals/test_terminate_all_agents.py

@@ -0,0 +1,51 @@
+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(build_flask_client):
+    container = StubDIContainer()
+
+    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:
+        yield flask_client
+
+
+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