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..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/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.") diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index c56e13322..86616b596 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, @@ -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): diff --git a/monkey/monkey_island/cc/resources/__init__.py b/monkey/monkey_island/cc/resources/__init__.py index b13c6cc06..937766e2b 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, TerminateAllAgents 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)} 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/agent_signals.py b/monkey/monkey_island/cc/resources/agent_signals/agent_signals.py new file mode 100644 index 000000000..f6ce22f71 --- /dev/null +++ b/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/"] + + 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 diff --git a/monkey/monkey_island/cc/resources/agent_signals/terminate_all_agents.py b/monkey/monkey_island/cc/resources/agent_signals/terminate_all_agents.py new file mode 100644 index 000000000..70032f11e --- /dev/null +++ b/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 diff --git a/monkey/monkey_island/cc/services/infection_lifecycle.py b/monkey/monkey_island/cc/services/infection_lifecycle.py index 937a3abeb..ef11ce5b1 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) @@ -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(): 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..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,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})}); 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 new file mode 100644 index 000000000..e7a20a9c9 --- /dev/null +++ b/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 diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/agent_signals/test_terminate_all_agents.py b/monkey/tests/unit_tests/monkey_island/cc/resources/agent_signals/test_terminate_all_agents.py new file mode 100644 index 000000000..bcf91afd4 --- /dev/null +++ b/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