From d084b003678b8f39742dabd15f66d92fa14db2ec Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 26 Apr 2022 02:36:17 -0400 Subject: [PATCH] Island: Use DIContainer to construct custom PBA resources --- monkey/monkey_island/cc/app.py | 109 ++++++++++-------- monkey/monkey_island/cc/server_setup.py | 15 ++- .../monkey_island/cc/services/initialize.py | 17 ++- 3 files changed, 86 insertions(+), 55 deletions(-) diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index beb9e1468..da8d312f7 100644 --- a/monkey/monkey_island/cc/app.py +++ b/monkey/monkey_island/cc/app.py @@ -1,12 +1,13 @@ import os import uuid from datetime import timedelta -from pathlib import Path +from typing import Type import flask_restful from flask import Flask, Response, send_from_directory from werkzeug.exceptions import NotFound +from common import DIContainer from monkey_island.cc.database import database, mongo from monkey_island.cc.resources.agent_controls import StopAgentCheck, StopAllAgents from monkey_island.cc.resources.attack.attack_report import AttackReport @@ -47,7 +48,6 @@ from monkey_island.cc.resources.zero_trust.finding_event import ZeroTrustFinding from monkey_island.cc.resources.zero_trust.zero_trust_report import ZeroTrustReport from monkey_island.cc.server_utils.consts import MONKEY_ISLAND_ABS_PATH from monkey_island.cc.server_utils.custom_json_encoder import CustomJSONEncoder -from monkey_island.cc.services import DirectoryFileStorageService from monkey_island.cc.services.remote_run_aws import RemoteRunAwsService from monkey_island.cc.services.representations import output_json @@ -113,73 +113,90 @@ def init_app_url_rules(app): app.add_url_rule("/", "serve_static_file", serve_static_file) -def init_api_resources(api, data_dir: Path): - api.add_resource(Root, "/api") - api.add_resource(Registration, "/api/registration") - api.add_resource(Authenticate, "/api/auth") - api.add_resource( +# TODO: Come up with a better name than FlaskResourceManager +class FlaskResourceManager: + def __init__(self, api: flask_restful.Api, container: DIContainer): + self._api = api + self._container = container + + def add_resource(self, resource: Type[flask_restful.Resource], *urls: str): + dependencies = self._container.resolve_dependencies(resource) + self._api.add_resource(resource, *urls, resource_class_args=dependencies) + + +def init_api_resources(flask_resource_manager: FlaskResourceManager): + flask_resource_manager.add_resource(Root, "/api") + flask_resource_manager.add_resource(Registration, "/api/registration") + flask_resource_manager.add_resource(Authenticate, "/api/auth") + flask_resource_manager.add_resource( Monkey, "/api/agent", "/api/agent/", "/api/agent//", ) - api.add_resource(LocalRun, "/api/local-monkey") - api.add_resource(Telemetry, "/api/telemetry", "/api/telemetry/") + flask_resource_manager.add_resource(LocalRun, "/api/local-monkey") + flask_resource_manager.add_resource( + Telemetry, "/api/telemetry", "/api/telemetry/" + ) - api.add_resource(IslandMode, "/api/island-mode") - api.add_resource(IslandConfiguration, "/api/configuration/island") - api.add_resource(ConfigurationExport, "/api/configuration/export") - api.add_resource(ConfigurationImport, "/api/configuration/import") - api.add_resource( + flask_resource_manager.add_resource(IslandMode, "/api/island-mode") + flask_resource_manager.add_resource(IslandConfiguration, "/api/configuration/island") + flask_resource_manager.add_resource(ConfigurationExport, "/api/configuration/export") + flask_resource_manager.add_resource(ConfigurationImport, "/api/configuration/import") + flask_resource_manager.add_resource( MonkeyDownload, "/api/agent/download/", ) - 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/node-states") + flask_resource_manager.add_resource(NetMap, "/api/netmap") + flask_resource_manager.add_resource(Edge, "/api/netmap/edge") + flask_resource_manager.add_resource(Node, "/api/netmap/node") + flask_resource_manager.add_resource(NodeStates, "/api/netmap/node-states") - api.add_resource(SecurityReport, "/api/report/security") - api.add_resource(ZeroTrustReport, "/api/report/zero-trust/") - api.add_resource(AttackReport, "/api/report/attack") - api.add_resource(RansomwareReport, "/api/report/ransomware") - api.add_resource(ManualExploitation, "/api/exploitations/manual") - api.add_resource(MonkeyExploitation, "/api/exploitations/monkey") + flask_resource_manager.add_resource(SecurityReport, "/api/report/security") + flask_resource_manager.add_resource( + ZeroTrustReport, "/api/report/zero-trust/" + ) + flask_resource_manager.add_resource(AttackReport, "/api/report/attack") + flask_resource_manager.add_resource(RansomwareReport, "/api/report/ransomware") + flask_resource_manager.add_resource(ManualExploitation, "/api/exploitations/manual") + flask_resource_manager.add_resource(MonkeyExploitation, "/api/exploitations/monkey") - api.add_resource(ZeroTrustFindingEvent, "/api/zero-trust/finding-event/") - api.add_resource(TelemetryFeed, "/api/telemetry-feed") - api.add_resource(Log, "/api/log") - api.add_resource(IslandLog, "/api/log/island/download") + flask_resource_manager.add_resource( + ZeroTrustFindingEvent, "/api/zero-trust/finding-event/" + ) + flask_resource_manager.add_resource(TelemetryFeed, "/api/telemetry-feed") + flask_resource_manager.add_resource(Log, "/api/log") + flask_resource_manager.add_resource(IslandLog, "/api/log/island/download") - # This is temporary until we get DI all worked out. - file_storage_service = DirectoryFileStorageService(data_dir / "custom_pbas") - api.add_resource( + flask_resource_manager.add_resource( PBAFileDownload, "/api/pba/download/", - resource_class_kwargs={"file_storage_service": file_storage_service}, ) - api.add_resource( + flask_resource_manager.add_resource( FileUpload, "/api/file-upload/", "/api/file-upload/?load=", "/api/file-upload/?restore=", - resource_class_kwargs={"file_storage_service": file_storage_service}, ) - api.add_resource(PropagationCredentials, "/api/propagation-credentials/") - api.add_resource(RemoteRun, "/api/remote-monkey") - api.add_resource(VersionUpdate, "/api/version-update") - api.add_resource(StopAgentCheck, "/api/monkey-control/needs-to-stop/") - api.add_resource(StopAllAgents, "/api/monkey-control/stop-all-agents") + flask_resource_manager.add_resource( + PropagationCredentials, "/api/propagation-credentials/" + ) + flask_resource_manager.add_resource(RemoteRun, "/api/remote-monkey") + flask_resource_manager.add_resource(VersionUpdate, "/api/version-update") + flask_resource_manager.add_resource( + StopAgentCheck, "/api/monkey-control/needs-to-stop/" + ) + flask_resource_manager.add_resource(StopAllAgents, "/api/monkey-control/stop-all-agents") # Resources used by black box tests - api.add_resource(MonkeyBlackboxEndpoint, "/api/test/monkey") - api.add_resource(ClearCaches, "/api/test/clear-caches") - api.add_resource(LogBlackboxEndpoint, "/api/test/log") - api.add_resource(TelemetryBlackboxEndpoint, "/api/test/telemetry") + flask_resource_manager.add_resource(MonkeyBlackboxEndpoint, "/api/test/monkey") + flask_resource_manager.add_resource(ClearCaches, "/api/test/clear-caches") + flask_resource_manager.add_resource(LogBlackboxEndpoint, "/api/test/log") + flask_resource_manager.add_resource(TelemetryBlackboxEndpoint, "/api/test/telemetry") -def init_app(mongo_url, data_dir: Path): +def init_app(mongo_url: str, container: DIContainer): app = Flask(__name__) api = flask_restful.Api(app) @@ -188,6 +205,8 @@ def init_app(mongo_url, data_dir: Path): init_app_config(app, mongo_url) init_app_services(app) init_app_url_rules(app) - init_api_resources(api, data_dir) + + flask_resource_manager = FlaskResourceManager(api, container) + init_api_resources(flask_resource_manager) return app diff --git a/monkey/monkey_island/cc/server_setup.py b/monkey/monkey_island/cc/server_setup.py index 072cb32a8..df41bb137 100644 --- a/monkey/monkey_island/cc/server_setup.py +++ b/monkey/monkey_island/cc/server_setup.py @@ -16,6 +16,7 @@ MONKEY_ISLAND_DIR_BASE_PATH = str(Path(__file__).parent.parent) if str(MONKEY_ISLAND_DIR_BASE_PATH) not in sys.path: sys.path.insert(0, MONKEY_ISLAND_DIR_BASE_PATH) +from common import DIContainer # noqa: E402 from common.version import get_version # noqa: E402 from monkey_island.cc.app import init_app # noqa: E402 from monkey_island.cc.arg_parser import IslandCmdArgs # noqa: E402 @@ -47,7 +48,7 @@ def run_monkey_island(): _exit_on_invalid_config_options(config_options) _configure_logging(config_options) - _initialize_globals(config_options.data_dir) + container = _initialize_globals(config_options.data_dir) mongo_db_process = None if config_options.start_mongodb: @@ -56,7 +57,7 @@ def run_monkey_island(): _connect_to_mongodb(mongo_db_process) _configure_gevent_exception_handling(config_options.data_dir) - _start_island_server(island_args.setup_only, config_options) + _start_island_server(island_args.setup_only, config_options, container) def _extract_config(island_args: IslandCmdArgs) -> IslandConfigOptions: @@ -88,8 +89,8 @@ def _configure_logging(config_options): setup_logging(config_options.data_dir, config_options.log_level) -def _initialize_globals(data_dir: Path): - initialize_services(data_dir) +def _initialize_globals(data_dir: Path) -> DIContainer: + return initialize_services(data_dir) def _start_mongodb(data_dir: Path) -> MongoDbProcess: @@ -127,9 +128,11 @@ def _configure_gevent_exception_handling(data_dir): hub.handle_error = GeventHubErrorHandler(hub, logger) -def _start_island_server(should_setup_only, config_options: IslandConfigOptions): +def _start_island_server( + should_setup_only: bool, config_options: IslandConfigOptions, container: DIContainer +): populate_exporter_list() - app = init_app(mongo_setup.MONGO_URL, Path(config_options.data_dir)) + app = init_app(mongo_setup.MONGO_URL, container) if should_setup_only: logger.warning("Setup only flag passed. Exiting.") diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index f3ff10395..4a4b2e4af 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -1,13 +1,22 @@ -from monkey_island.cc.services import DirectoryFileStorageService +from pathlib import Path + +from common import DIContainer +from monkey_island.cc.services import DirectoryFileStorageService, IFileStorageService from monkey_island.cc.services.post_breach_files import PostBreachFilesService from monkey_island.cc.services.run_local_monkey import LocalMonkeyRunService from . import AuthenticationService, JsonFileUserDatastore -def initialize_services(data_dir): - # This is temporary until we get DI all worked out. - PostBreachFilesService.initialize(DirectoryFileStorageService(data_dir / "custom_pbas")) +def initialize_services(data_dir: Path) -> DIContainer: + container = DIContainer() + container.register_instance( + IFileStorageService, DirectoryFileStorageService(data_dir / "custom_pbas") + ) + # This is temporary until we get DI all worked out. + PostBreachFilesService.initialize(container.resolve(IFileStorageService)) LocalMonkeyRunService.initialize(data_dir) AuthenticationService.initialize(data_dir, JsonFileUserDatastore(data_dir)) + + return container