From 4fdd3370cafdbe027deb7068070d4b466c48617b Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Fri, 3 Dec 2021 16:21:31 +0200 Subject: [PATCH] Island, UI: implement the endpoint for stopping all monkeys, change the UI to call this endpoint and send a timestamp of button press --- .../master/control_channel.py | 5 +- monkey/monkey_island/cc/app.py | 32 +++---- .../cc/models/agent_controls/__init__.py | 1 + .../models/agent_controls/agent_controls.py | 7 ++ monkey/monkey_island/cc/models/config.py | 3 +- .../agent_controls/stop_agent_check.py | 7 +- .../agent_controls/stop_all_agents.py | 18 ++++ monkey/monkey_island/cc/resources/root.py | 6 +- monkey/monkey_island/cc/services/database.py | 2 + .../cc/services/infection_lifecycle.py | 94 +++++++++++-------- .../map/preview-pane/PreviewPane.js | 1 + .../cc/ui/src/components/pages/MapPage.js | 9 +- vulture_allowlist.py | 3 + 13 files changed, 121 insertions(+), 67 deletions(-) create mode 100644 monkey/monkey_island/cc/models/agent_controls/__init__.py create mode 100644 monkey/monkey_island/cc/models/agent_controls/agent_controls.py create mode 100644 monkey/monkey_island/cc/resources/agent_controls/stop_all_agents.py diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py index 3509cedc2..17a2d3287 100644 --- a/monkey/infection_monkey/master/control_channel.py +++ b/monkey/infection_monkey/master/control_channel.py @@ -19,9 +19,12 @@ class ControlChannel(IControlChannel): self._control_channel_server = server def should_agent_stop(self) -> bool: + if not self._control_channel_server: + logger.error("Agent should stop because it can't connect to the C&C server.") + return True try: response = requests.get( # noqa: DUO123 - f"{self._control_channel_server}/api/monkey_control/{self._agent_id}", + f"https://{self._control_channel_server}/api/monkey_control/needs-to-stop/{self._agent_id}", verify=False, proxies=ControlClient.proxies, timeout=SHORT_REQUEST_TIMEOUT, diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index d3628c475..ea0556720 100644 --- a/monkey/monkey_island/cc/app.py +++ b/monkey/monkey_island/cc/app.py @@ -11,6 +11,7 @@ from monkey_island.cc.database import database, mongo from monkey_island.cc.resources.agent_controls import ( StartedOnIsland, StopAgentCheck, + StopAllAgents, ) from monkey_island.cc.resources.attack.attack_report import AttackReport from monkey_island.cc.resources.auth.auth import Authenticate, init_jwt @@ -131,25 +132,23 @@ def init_api_resources(api): "/api/monkey//", ) api.add_resource(Bootloader, "/api/bootloader/") - api.add_resource(LocalRun, "/api/local-monkey", "/api/local-monkey/") - api.add_resource(ClientRun, "/api/client-monkey", "/api/client-monkey/") - api.add_resource( - Telemetry, "/api/telemetry", "/api/telemetry/", "/api/telemetry/" - ) + api.add_resource(LocalRun, "/api/local-monkey") + api.add_resource(StopAgentCheck, "/api/local-monkey") + api.add_resource(ClientRun, "/api/client-monkey") + api.add_resource(Telemetry, "/api/telemetry", "/api/telemetry/") api.add_resource(IslandMode, "/api/island-mode") - api.add_resource(IslandConfiguration, "/api/configuration/island", "/api/configuration/island/") + api.add_resource(IslandConfiguration, "/api/configuration/island") api.add_resource(ConfigurationExport, "/api/configuration/export") api.add_resource(ConfigurationImport, "/api/configuration/import") api.add_resource( MonkeyDownload, "/api/monkey/download", - "/api/monkey/download/", "/api/monkey/download/", ) - api.add_resource(NetMap, "/api/netmap", "/api/netmap/") - api.add_resource(Edge, "/api/netmap/edge", "/api/netmap/edge/") - api.add_resource(Node, "/api/netmap/node", "/api/netmap/node/") + api.add_resource(NetMap, "/api/netmap") + api.add_resource(Edge, "/api/netmap/edge") + api.add_resource(Node, "/api/netmap/node") api.add_resource(NodeStates, "/api/netmap/nodeStates") api.add_resource(SecurityReport, "/api/report/security") @@ -160,9 +159,9 @@ def init_api_resources(api): api.add_resource(MonkeyExploitation, "/api/exploitations/monkey") api.add_resource(ZeroTrustFindingEvent, "/api/zero-trust/finding-event/") - api.add_resource(TelemetryFeed, "/api/telemetry-feed", "/api/telemetry-feed/") - api.add_resource(Log, "/api/log", "/api/log/") - api.add_resource(IslandLog, "/api/log/island/download", "/api/log/island/download/") + api.add_resource(TelemetryFeed, "/api/telemetry-feed") + api.add_resource(Log, "/api/log") + api.add_resource(IslandLog, "/api/log/island/download") api.add_resource(PBAFileDownload, "/api/pba/download/") api.add_resource(T1216PBAFileDownload, T1216_PBA_FILE_DOWNLOAD_PATH) api.add_resource( @@ -172,10 +171,11 @@ def init_api_resources(api): "/api/fileUpload/?restore=", ) api.add_resource(PropagationCredentials, "/api/propagation-credentials/") - api.add_resource(RemoteRun, "/api/remote-monkey", "/api/remote-monkey/") - api.add_resource(VersionUpdate, "/api/version-update", "/api/version-update/") + api.add_resource(RemoteRun, "/api/remote-monkey") + api.add_resource(VersionUpdate, "/api/version-update") api.add_resource(StartedOnIsland, "/api/monkey_control/started_on_island") - api.add_resource(StopAgentCheck, "/api/monkey_control/") + api.add_resource(StopAgentCheck, "/api/monkey_control/needs-to-stop/") + api.add_resource(StopAllAgents, "/api/monkey_control/stop-all-agents") api.add_resource(ScoutSuiteAuth, "/api/scoutsuite_auth/") api.add_resource(AWSKeys, "/api/aws_keys") diff --git a/monkey/monkey_island/cc/models/agent_controls/__init__.py b/monkey/monkey_island/cc/models/agent_controls/__init__.py new file mode 100644 index 000000000..e623955c3 --- /dev/null +++ b/monkey/monkey_island/cc/models/agent_controls/__init__.py @@ -0,0 +1 @@ +from .agent_controls import AgentControls diff --git a/monkey/monkey_island/cc/models/agent_controls/agent_controls.py b/monkey/monkey_island/cc/models/agent_controls/agent_controls.py new file mode 100644 index 000000000..37903d5e7 --- /dev/null +++ b/monkey/monkey_island/cc/models/agent_controls/agent_controls.py @@ -0,0 +1,7 @@ +from mongoengine import Document, FloatField + + +class AgentControls(Document): + + # Timestamp of the last "kill all agents" command + last_stop_all = FloatField(default=None) diff --git a/monkey/monkey_island/cc/models/config.py b/monkey/monkey_island/cc/models/config.py index f4af7b400..f2b82a8b4 100644 --- a/monkey/monkey_island/cc/models/config.py +++ b/monkey/monkey_island/cc/models/config.py @@ -1,4 +1,4 @@ -from mongoengine import EmbeddedDocument +from mongoengine import EmbeddedDocument, BooleanField class Config(EmbeddedDocument): @@ -8,5 +8,6 @@ class Config(EmbeddedDocument): See https://mongoengine-odm.readthedocs.io/apireference.html#mongoengine.FieldDoesNotExist """ + alive = BooleanField() meta = {"strict": False} pass diff --git a/monkey/monkey_island/cc/resources/agent_controls/stop_agent_check.py b/monkey/monkey_island/cc/resources/agent_controls/stop_agent_check.py index 817d6db94..3fb948a68 100644 --- a/monkey/monkey_island/cc/resources/agent_controls/stop_agent_check.py +++ b/monkey/monkey_island/cc/resources/agent_controls/stop_agent_check.py @@ -1,9 +1,8 @@ import flask_restful +from monkey_island.cc.services.infection_lifecycle import should_agent_die + class StopAgentCheck(flask_restful.Resource): def get(self, monkey_guid: int): - if monkey_guid % 2: - return {"stop_agent": True} - else: - return {"stop_agent": False} + return {"stop_agent": should_agent_die(monkey_guid)} 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 new file mode 100644 index 000000000..8d9c558ad --- /dev/null +++ b/monkey/monkey_island/cc/resources/agent_controls/stop_all_agents.py @@ -0,0 +1,18 @@ +import json + +import flask_restful +from flask import make_response, request + +from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.services.infection_lifecycle import set_stop_all + + +class StopAllAgents(flask_restful.Resource): + @jwt_required + def post(self): + data = json.loads(request.data) + if data["kill_time"]: + set_stop_all(data["kill_time"]) + return make_response({}, 200) + else: + return make_response({}, 400) diff --git a/monkey/monkey_island/cc/resources/root.py b/monkey/monkey_island/cc/resources/root.py index 41ff4e3ad..d3a36e6a2 100644 --- a/monkey/monkey_island/cc/resources/root.py +++ b/monkey/monkey_island/cc/resources/root.py @@ -6,7 +6,7 @@ from flask import jsonify, make_response, request from monkey_island.cc.database import mongo from monkey_island.cc.resources.auth.auth import jwt_required from monkey_island.cc.services.database import Database -from monkey_island.cc.services.infection_lifecycle import InfectionLifecycle +from monkey_island.cc.services.infection_lifecycle import get_completed_steps from monkey_island.cc.services.utils.network_utils import local_ip_addresses logger = logging.getLogger(__name__) @@ -21,8 +21,6 @@ class Root(flask_restful.Resource): return self.get_server_info() elif action == "reset": return jwt_required(Database.reset_db)() - elif action == "killall": - return jwt_required(InfectionLifecycle.kill_all)() elif action == "is-up": return {"is-up": True} else: @@ -33,5 +31,5 @@ class Root(flask_restful.Resource): return jsonify( ip_addresses=local_ip_addresses(), mongo=str(mongo.db), - completed_steps=InfectionLifecycle.get_completed_steps(), + completed_steps=get_completed_steps(), ) diff --git a/monkey/monkey_island/cc/services/database.py b/monkey/monkey_island/cc/services/database.py index 027bd49e2..14f7296f4 100644 --- a/monkey/monkey_island/cc/services/database.py +++ b/monkey/monkey_island/cc/services/database.py @@ -4,6 +4,7 @@ from flask import jsonify from monkey_island.cc.database import mongo from monkey_island.cc.models.attack.attack_mitigations import AttackMitigations +from monkey_island.cc.services.infection_lifecycle import init_agent_controls from monkey_island.cc.services.config import ConfigService logger = logging.getLogger(__name__) @@ -23,6 +24,7 @@ class Database(object): if not x.startswith("system.") and not x == AttackMitigations.COLLECTION_NAME ] ConfigService.init_config() + init_agent_controls() logger.info("DB was reset") return jsonify(status="OK") diff --git a/monkey/monkey_island/cc/services/infection_lifecycle.py b/monkey/monkey_island/cc/services/infection_lifecycle.py index 5529cc70d..55c9f79b9 100644 --- a/monkey/monkey_island/cc/services/infection_lifecycle.py +++ b/monkey/monkey_island/cc/services/infection_lifecycle.py @@ -1,9 +1,7 @@ import logging -from datetime import datetime -from flask import jsonify - -from monkey_island.cc.database import mongo +from monkey_island.cc.models import Monkey +from monkey_island.cc.models.agent_controls import AgentControls from monkey_island.cc.resources.blackbox.utils.telem_store import TestTelemStore from monkey_island.cc.services.config import ConfigService from monkey_island.cc.services.node import NodeService @@ -16,42 +14,60 @@ from monkey_island.cc.services.reporting.report_generation_synchronisation impor logger = logging.getLogger(__name__) -class InfectionLifecycle: - @staticmethod - def kill_all(): - mongo.db.monkey.update( - {"dead": False}, - {"$set": {"config.alive": False, "modifytime": datetime.now()}}, - upsert=False, - multi=True, - ) - logger.info("Kill all monkeys was called") - return jsonify(status="OK") +def set_stop_all(time: float): + for monkey in Monkey.objects(): + monkey.config.alive = False + monkey.save() + agent_controls = AgentControls.objects.first() + agent_controls.last_stop_all = time + agent_controls.save() - @staticmethod - def get_completed_steps(): - is_any_exists = NodeService.is_any_monkey_exists() - infection_done = NodeService.is_monkey_finished_running() - if infection_done: - InfectionLifecycle._on_finished_infection() - report_done = ReportService.is_report_generated() - else: # Infection is not done - report_done = False +def should_agent_die(guid: int) -> bool: + monkey = Monkey.objects(guid=str(guid)).first() + return _is_monkey_marked_dead(monkey) or _is_monkey_killed_manually(monkey) - return dict( - run_server=True, - run_monkey=is_any_exists, - infection_done=infection_done, - report_done=report_done, - ) - @staticmethod - def _on_finished_infection(): - # Checking is_report_being_generated here, because we don't want to wait to generate a - # report; rather, - # we want to skip and reply. - if not is_report_being_generated() and not ReportService.is_latest_report_exists(): - safe_generate_reports() - if ConfigService.is_test_telem_export_enabled() and not TestTelemStore.TELEMS_EXPORTED: - TestTelemStore.export_telems() +def _is_monkey_marked_dead(monkey: Monkey) -> bool: + return monkey.config.alive + + +def _is_monkey_killed_manually(monkey: Monkey) -> bool: + if monkey.has_parent(): + launch_timestamp = monkey.get_parent().launch_time + else: + launch_timestamp = monkey.launch_time + kill_timestamp = AgentControls.objects.first().last_stop_all + return int(kill_timestamp) >= int(launch_timestamp) + + +def init_agent_controls(): + AgentControls().save() + + +def get_completed_steps(): + is_any_exists = NodeService.is_any_monkey_exists() + infection_done = NodeService.is_monkey_finished_running() + + if infection_done: + _on_finished_infection() + report_done = ReportService.is_report_generated() + else: # Infection is not done + report_done = False + + return dict( + run_server=True, + run_monkey=is_any_exists, + infection_done=infection_done, + report_done=report_done, + ) + + +def _on_finished_infection(): + # Checking is_report_being_generated here, because we don't want to wait to generate a + # report; rather, + # we want to skip and reply. + if not is_report_being_generated() and not ReportService.is_latest_report_exists(): + safe_generate_reports() + if ConfigService.is_test_telem_export_enabled() and not TestTelemStore.TELEMS_EXPORTED: + TestTelemStore.export_telems() diff --git a/monkey/monkey_island/cc/ui/src/components/map/preview-pane/PreviewPane.js b/monkey/monkey_island/cc/ui/src/components/map/preview-pane/PreviewPane.js index 9007194b0..81e1d3c9d 100644 --- a/monkey/monkey_island/cc/ui/src/components/map/preview-pane/PreviewPane.js +++ b/monkey/monkey_island/cc/ui/src/components/map/preview-pane/PreviewPane.js @@ -78,6 +78,7 @@ class PreviewPaneComponent extends AuthComponent { }); } + // TODO remove this forceKillRow(asset) { return ( 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 6026cebb6..8dcfe0ce6 100644 --- a/monkey/monkey_island/cc/ui/src/components/pages/MapPage.js +++ b/monkey/monkey_island/cc/ui/src/components/pages/MapPage.js @@ -84,9 +84,14 @@ class MapPageComponent extends AuthComponent { } killAllMonkeys = () => { - this.authFetch('/api?action=killall') + this.authFetch('/api/agent_control/stop-all-agents', + { + method: 'POST', + headers: {'Content-Type': 'application/json'}, + body: JSON.stringify({kill_time: Date.now()}) + }) .then(res => res.json()) - .then(res => this.setState({killPressed: (res.status === 'OK')})); + .then(res => {this.setState({killPressed: true}); console.log(res)}); }; renderKillDialogModal = () => { diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 20c130c33..7c9917984 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -3,6 +3,7 @@ Everything in this file is what Vulture found as dead code but either isn't real dead or is kept deliberately. Referencing these in a file like this makes sure that Vulture doesn't mark these as dead again. """ +from monkey_island.cc import app from monkey_island.cc.models import Report fake_monkey_dir_path # unused variable (monkey/tests/infection_monkey/post_breach/actions/test_users_custom_pba.py:37) @@ -100,6 +101,8 @@ EnvironmentCollector # unused class (monkey/infection_monkey/system_info/collec ProcessListCollector # unused class (monkey/infection_monkey/system_info/collectors/process_list_collector.py:18) _.coinit_flags # unused attribute (monkey/infection_monkey/system_info/windows_info_collector.py:11) _.representations # unused attribute (monkey/monkey_island/cc/app.py:180) +_.representations # unused attribute (monkey/monkey_island/cc/app.py:180) +app.url_map.strict_slashes _.log_message # unused method (monkey/infection_monkey/transport/http.py:188) _.log_message # unused method (monkey/infection_monkey/transport/http.py:109) _.version_string # unused method (monkey/infection_monkey/transport/http.py:148)