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

2261 add agent signals resource
This commit is contained in:
Mike Salvatore 2022-09-23 09:19:33 -04:00
commit 6ae0e6f715
12 changed files with 187 additions and 46 deletions

View File

@ -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.")

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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