diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index 07e08ea9a..f7a7a3e01 100644 --- a/monkey/monkey_island/cc/app.py +++ b/monkey/monkey_island/cc/app.py @@ -1,7 +1,8 @@ import os +import re import uuid from datetime import timedelta -from typing import Type +from typing import Iterable, Type import flask_restful from flask import Flask, Response, send_from_directory @@ -10,6 +11,7 @@ from werkzeug.exceptions import NotFound from common import DIContainer from monkey_island.cc.database import database, mongo from monkey_island.cc.resources import RemoteRun +from monkey_island.cc.resources.AbstractResource import AbstractResource from monkey_island.cc.resources.agent_controls import StopAgentCheck, StopAllAgents from monkey_island.cc.resources.attack.attack_report import AttackReport from monkey_island.cc.resources.auth.auth import Authenticate, init_jwt @@ -109,75 +111,82 @@ def init_app_url_rules(app): class FlaskDIWrapper: + class DuplicateURLError(Exception): + pass + + url_parameter_regex = re.compile(r"<.*?:.*?>") + def __init__(self, api: flask_restful.Api, container: DIContainer): self._api = api self._container = container + self._reserved_urls = set() + + def add_resource(self, resource: Type[AbstractResource]): + if len(resource.urls) == 0: + raise ValueError(f"Resource {resource.__name__} has no defined URLs") + + self._reserve_urls(resource.urls) - 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) + self._api.add_resource(resource, *resource.urls, resource_class_args=dependencies) + + def _reserve_urls(self, urls: Iterable[str]): + for url in map(FlaskDIWrapper._format_url, urls): + if url in self._reserved_urls: + raise FlaskDIWrapper.DuplicateURLError(f"URL {url} has already been registered!") + + self._reserved_urls.add(url) + + @staticmethod + def _format_url(url: str): + new_url = url.strip("/") + return FlaskDIWrapper.url_parameter_regex.sub("", new_url) def init_api_resources(api: FlaskDIWrapper): - api.add_resource(Root, "/api") - api.add_resource(Registration, "/api/registration") - api.add_resource(Authenticate, "/api/auth") - api.add_resource( - Monkey, - "/api/agent", - "/api/agent/", - "/api/agent//", - ) - api.add_resource(LocalRun, "/api/local-monkey") - api.add_resource(Telemetry, "/api/telemetry", "/api/telemetry/") + api.add_resource(Root) + api.add_resource(Registration) + api.add_resource(Authenticate) + api.add_resource(Monkey) + api.add_resource(LocalRun) + api.add_resource(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( - 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") + api.add_resource(IslandMode) + api.add_resource(IslandConfiguration) + api.add_resource(ConfigurationExport) + api.add_resource(ConfigurationImport) + api.add_resource(MonkeyDownload) + api.add_resource(NetMap) + api.add_resource(Edge) + api.add_resource(Node) + api.add_resource(NodeStates) - 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") + api.add_resource(SecurityReport) + api.add_resource(ZeroTrustReport) + api.add_resource(AttackReport) + api.add_resource(RansomwareReport) + api.add_resource(ManualExploitation) + api.add_resource(MonkeyExploitation) - 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") + api.add_resource(ZeroTrustFindingEvent) + api.add_resource(TelemetryFeed) + api.add_resource(Log) + api.add_resource(IslandLog) - api.add_resource( - PBAFileDownload, - "/api/pba/download/", - ) - api.add_resource( - FileUpload, - "/api/file-upload/", - "/api/file-upload/?load=", - "/api/file-upload/?restore=", - ) + api.add_resource(PBAFileDownload) + api.add_resource(FileUpload) - 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") + api.add_resource(PropagationCredentials) + api.add_resource(RemoteRun) + api.add_resource(VersionUpdate) + api.add_resource(StopAgentCheck) + api.add_resource(StopAllAgents) # 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") + api.add_resource(MonkeyBlackboxEndpoint) + api.add_resource(ClearCaches) + api.add_resource(LogBlackboxEndpoint) + api.add_resource(TelemetryBlackboxEndpoint) def init_app(mongo_url: str, container: DIContainer): diff --git a/monkey/monkey_island/cc/resources/AbstractResource.py b/monkey/monkey_island/cc/resources/AbstractResource.py new file mode 100644 index 000000000..799cd5a23 --- /dev/null +++ b/monkey/monkey_island/cc/resources/AbstractResource.py @@ -0,0 +1,6 @@ +import flask_restful + + +# The purpose of this class is to decouple resources from flask +class AbstractResource(flask_restful.Resource): + urls = [] 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 3fb948a68..e26b2d172 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,8 +1,9 @@ -import flask_restful - +from monkey_island.cc.resources.AbstractResource import AbstractResource from monkey_island.cc.services.infection_lifecycle import should_agent_die -class StopAgentCheck(flask_restful.Resource): +class StopAgentCheck(AbstractResource): + urls = ["/api/monkey-control/needs-to-stop/"] + def get(self, monkey_guid: int): 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 index a8819243b..82667ecf2 100644 --- a/monkey/monkey_island/cc/resources/agent_controls/stop_all_agents.py +++ b/monkey/monkey_island/cc/resources/agent_controls/stop_all_agents.py @@ -1,14 +1,16 @@ 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.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(flask_restful.Resource): +class StopAllAgents(AbstractResource): + urls = ["/api/monkey-control/stop-all-agents"] + @jwt_required def post(self): with agent_killing_mutex: diff --git a/monkey/monkey_island/cc/resources/attack/attack_report.py b/monkey/monkey_island/cc/resources/attack/attack_report.py index 502538990..2dffa929d 100644 --- a/monkey/monkey_island/cc/resources/attack/attack_report.py +++ b/monkey/monkey_island/cc/resources/attack/attack_report.py @@ -1,12 +1,14 @@ -import flask_restful from flask import current_app, json -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services.attack.attack_report import AttackReportService from monkey_island.cc.services.attack.attack_schema import SCHEMA -class AttackReport(flask_restful.Resource): +class AttackReport(AbstractResource): + urls = ["/api/report/attack"] + @jwt_required def get(self): response_content = { diff --git a/monkey/monkey_island/cc/resources/auth/auth.py b/monkey/monkey_island/cc/resources/auth/auth.py index f5b73e062..57c78cfe0 100644 --- a/monkey/monkey_island/cc/resources/auth/auth.py +++ b/monkey/monkey_island/cc/resources/auth/auth.py @@ -1,14 +1,12 @@ import logging -from functools import wraps import flask_jwt_extended -import flask_restful from flask import make_response, request -from flask_jwt_extended.exceptions import JWTExtendedException -from jwt import PyJWTError from common.utils.exceptions import IncorrectCredentialsError +from monkey_island.cc.resources.AbstractResource import AbstractResource from monkey_island.cc.resources.auth.credential_utils import get_username_password_from_request +from monkey_island.cc.resources.request_authentication import create_access_token from monkey_island.cc.services import AuthenticationService logger = logging.getLogger(__name__) @@ -21,13 +19,15 @@ def init_jwt(app): ) -class Authenticate(flask_restful.Resource): +class Authenticate(AbstractResource): """ Resource for user authentication. The user provides the username and password and we give them a JWT. See `AuthService.js` file for the frontend counterpart for this code. """ + urls = ["/api/auth"] + def post(self): """ Example request: @@ -40,30 +40,8 @@ class Authenticate(flask_restful.Resource): try: AuthenticationService.authenticate(username, password) - access_token = _create_access_token(username) + access_token = create_access_token(username) except IncorrectCredentialsError: return make_response({"error": "Invalid credentials"}, 401) return make_response({"access_token": access_token, "error": ""}, 200) - - -def _create_access_token(username): - access_token = flask_jwt_extended.create_access_token(identity=username) - logger.debug(f"Created access token for user {username} that begins with {access_token[:4]}") - - return access_token - - -# See https://flask-jwt-extended.readthedocs.io/en/stable/custom_decorators/ -def jwt_required(fn): - @wraps(fn) - def wrapper(*args, **kwargs): - try: - flask_jwt_extended.verify_jwt_in_request() - return fn(*args, **kwargs) - # Catch authentication related errors in the verification or inside the called function. - # All other exceptions propagate - except (JWTExtendedException, PyJWTError) as e: - return make_response({"error": f"Authentication error: {str(e)}"}, 401) - - return wrapper diff --git a/monkey/monkey_island/cc/resources/auth/registration.py b/monkey/monkey_island/cc/resources/auth/registration.py index 175582733..1625de62b 100644 --- a/monkey/monkey_island/cc/resources/auth/registration.py +++ b/monkey/monkey_island/cc/resources/auth/registration.py @@ -1,16 +1,19 @@ import logging -import flask_restful from flask import make_response, request from common.utils.exceptions import AlreadyRegisteredError, InvalidRegistrationCredentialsError +from monkey_island.cc.resources.AbstractResource import AbstractResource from monkey_island.cc.resources.auth.credential_utils import get_username_password_from_request from monkey_island.cc.services import AuthenticationService logger = logging.getLogger(__name__) -class Registration(flask_restful.Resource): +class Registration(AbstractResource): + + urls = ["/api/registration"] + def get(self): return {"needs_registration": AuthenticationService.needs_registration()} diff --git a/monkey/monkey_island/cc/resources/blackbox/clear_caches.py b/monkey/monkey_island/cc/resources/blackbox/clear_caches.py index b8ebeb056..116bc0f07 100644 --- a/monkey/monkey_island/cc/resources/blackbox/clear_caches.py +++ b/monkey/monkey_island/cc/resources/blackbox/clear_caches.py @@ -2,7 +2,8 @@ import logging import flask_restful -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services.attack.attack_report import AttackReportService from monkey_island.cc.services.reporting.report import ReportService @@ -11,7 +12,8 @@ NOT_ALL_REPORTS_DELETED = "Not all reports have been cleared from the DB!" logger = logging.getLogger(__name__) -class ClearCaches(flask_restful.Resource): +class ClearCaches(AbstractResource): + urls = ["/api/test/clear-caches"] """ Used for timing tests - we want to get actual execution time of functions in BlackBox without caching - diff --git a/monkey/monkey_island/cc/resources/blackbox/log_blackbox_endpoint.py b/monkey/monkey_island/cc/resources/blackbox/log_blackbox_endpoint.py index c101b567a..a2700cc37 100644 --- a/monkey/monkey_island/cc/resources/blackbox/log_blackbox_endpoint.py +++ b/monkey/monkey_island/cc/resources/blackbox/log_blackbox_endpoint.py @@ -1,12 +1,14 @@ -import flask_restful from bson import json_util from flask import request from monkey_island.cc.database import database, mongo -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required -class LogBlackboxEndpoint(flask_restful.Resource): +class LogBlackboxEndpoint(AbstractResource): + urls = ["/api/test/log"] + @jwt_required def get(self): find_query = json_util.loads(request.args.get("find_query")) diff --git a/monkey/monkey_island/cc/resources/blackbox/monkey_blackbox_endpoint.py b/monkey/monkey_island/cc/resources/blackbox/monkey_blackbox_endpoint.py index 2957fd4b9..b02cf224f 100644 --- a/monkey/monkey_island/cc/resources/blackbox/monkey_blackbox_endpoint.py +++ b/monkey/monkey_island/cc/resources/blackbox/monkey_blackbox_endpoint.py @@ -1,12 +1,14 @@ -import flask_restful from bson import json_util from flask import request from monkey_island.cc.database import mongo -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required -class MonkeyBlackboxEndpoint(flask_restful.Resource): +class MonkeyBlackboxEndpoint(AbstractResource): + urls = ["/api/test/monkey"] + @jwt_required def get(self, **kw): find_query = json_util.loads(request.args.get("find_query")) diff --git a/monkey/monkey_island/cc/resources/blackbox/telemetry_blackbox_endpoint.py b/monkey/monkey_island/cc/resources/blackbox/telemetry_blackbox_endpoint.py index f1e958e3e..e3b91cc46 100644 --- a/monkey/monkey_island/cc/resources/blackbox/telemetry_blackbox_endpoint.py +++ b/monkey/monkey_island/cc/resources/blackbox/telemetry_blackbox_endpoint.py @@ -1,12 +1,14 @@ -import flask_restful from bson import json_util from flask import request from monkey_island.cc.models.telemetries import get_telemetry_by_query -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required -class TelemetryBlackboxEndpoint(flask_restful.Resource): +class TelemetryBlackboxEndpoint(AbstractResource): + urls = ["/api/test/telemetry"] + @jwt_required def get(self, **kw): find_query = json_util.loads(request.args.get("find_query")) diff --git a/monkey/monkey_island/cc/resources/configuration_export.py b/monkey/monkey_island/cc/resources/configuration_export.py index 111cfa177..e10e5d1a3 100644 --- a/monkey/monkey_island/cc/resources/configuration_export.py +++ b/monkey/monkey_island/cc/resources/configuration_export.py @@ -1,14 +1,16 @@ import json -import flask_restful from flask import request -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.server_utils.encryption import PasswordBasedStringEncryptor from monkey_island.cc.services.config import ConfigService -class ConfigurationExport(flask_restful.Resource): +class ConfigurationExport(AbstractResource): + urls = ["/api/configuration/export"] + @jwt_required def post(self): data = json.loads(request.data) diff --git a/monkey/monkey_island/cc/resources/configuration_import.py b/monkey/monkey_island/cc/resources/configuration_import.py index 3a66a2ed0..ecf94e986 100644 --- a/monkey/monkey_island/cc/resources/configuration_import.py +++ b/monkey/monkey_island/cc/resources/configuration_import.py @@ -3,11 +3,11 @@ import logging from dataclasses import dataclass from json.decoder import JSONDecodeError -import flask_restful from flask import request from common.utils.exceptions import InvalidConfigurationError -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.server_utils.encryption import ( InvalidCiphertextError, InvalidCredentialsError, @@ -38,7 +38,8 @@ class ResponseContents: return self.__dict__ -class ConfigurationImport(flask_restful.Resource): +class ConfigurationImport(AbstractResource): + urls = ["/api/configuration/import"] SUCCESS = False @jwt_required diff --git a/monkey/monkey_island/cc/resources/edge.py b/monkey/monkey_island/cc/resources/edge.py index 9eb0d5943..9798e116a 100644 --- a/monkey/monkey_island/cc/resources/edge.py +++ b/monkey/monkey_island/cc/resources/edge.py @@ -1,10 +1,12 @@ -import flask_restful from flask import request +from monkey_island.cc.resources.AbstractResource import AbstractResource from monkey_island.cc.services.edge.displayed_edge import DisplayedEdgeService -class Edge(flask_restful.Resource): +class Edge(AbstractResource): + urls = ["/api/netmap/edge"] + def get(self): edge_id = request.args.get("id") displayed_edge = DisplayedEdgeService.get_displayed_edge_by_id(edge_id) diff --git a/monkey/monkey_island/cc/resources/exploitations/manual_exploitation.py b/monkey/monkey_island/cc/resources/exploitations/manual_exploitation.py index 7c5db2f75..8006e31af 100644 --- a/monkey/monkey_island/cc/resources/exploitations/manual_exploitation.py +++ b/monkey/monkey_island/cc/resources/exploitations/manual_exploitation.py @@ -1,12 +1,13 @@ -import flask_restful - -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services.reporting.exploitations.manual_exploitation import ( get_manual_exploitations, ) -class ManualExploitation(flask_restful.Resource): +class ManualExploitation(AbstractResource): + urls = ["/api/exploitations/manual"] + @jwt_required def get(self): manual_exploitations = [ diff --git a/monkey/monkey_island/cc/resources/exploitations/monkey_exploitation.py b/monkey/monkey_island/cc/resources/exploitations/monkey_exploitation.py index 5e00a51a0..5cbf42a35 100644 --- a/monkey/monkey_island/cc/resources/exploitations/monkey_exploitation.py +++ b/monkey/monkey_island/cc/resources/exploitations/monkey_exploitation.py @@ -1,12 +1,13 @@ -import flask_restful - -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services.reporting.exploitations.monkey_exploitation import ( get_monkey_exploited, ) -class MonkeyExploitation(flask_restful.Resource): +class MonkeyExploitation(AbstractResource): + urls = ["/api/exploitations/monkey"] + @jwt_required def get(self): monkey_exploitations = [exploitation.__dict__ for exploitation in get_monkey_exploited()] diff --git a/monkey/monkey_island/cc/resources/island_configuration.py b/monkey/monkey_island/cc/resources/island_configuration.py index 42730e477..21da243ed 100644 --- a/monkey/monkey_island/cc/resources/island_configuration.py +++ b/monkey/monkey_island/cc/resources/island_configuration.py @@ -1,13 +1,16 @@ import json -import flask_restful from flask import abort, jsonify, request -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services.config import ConfigService -class IslandConfiguration(flask_restful.Resource): +class IslandConfiguration(AbstractResource): + + urls = ["/api/configuration/island"] + @jwt_required def get(self): return jsonify( diff --git a/monkey/monkey_island/cc/resources/island_logs.py b/monkey/monkey_island/cc/resources/island_logs.py index ae5bb1398..ec36b7a88 100644 --- a/monkey/monkey_island/cc/resources/island_logs.py +++ b/monkey/monkey_island/cc/resources/island_logs.py @@ -1,14 +1,15 @@ import logging -import flask_restful - -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services.island_logs import IslandLogService logger = logging.getLogger(__name__) -class IslandLog(flask_restful.Resource): +class IslandLog(AbstractResource): + urls = ["/api/log/island/download"] + @jwt_required def get(self): try: diff --git a/monkey/monkey_island/cc/resources/island_mode.py b/monkey/monkey_island/cc/resources/island_mode.py index 389d79dea..fed0d0aff 100644 --- a/monkey/monkey_island/cc/resources/island_mode.py +++ b/monkey/monkey_island/cc/resources/island_mode.py @@ -1,10 +1,10 @@ import json import logging -import flask_restful from flask import make_response, request -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services.config_manipulator import update_config_on_mode_set from monkey_island.cc.services.mode.island_mode_service import ModeNotSetError, get_mode, set_mode from monkey_island.cc.services.mode.mode_enum import IslandModeEnum @@ -12,7 +12,9 @@ from monkey_island.cc.services.mode.mode_enum import IslandModeEnum logger = logging.getLogger(__name__) -class IslandMode(flask_restful.Resource): +class IslandMode(AbstractResource): + urls = ["/api/island-mode"] + @jwt_required def post(self): try: diff --git a/monkey/monkey_island/cc/resources/local_run.py b/monkey/monkey_island/cc/resources/local_run.py index 5645557da..7de0b844a 100644 --- a/monkey/monkey_island/cc/resources/local_run.py +++ b/monkey/monkey_island/cc/resources/local_run.py @@ -1,15 +1,18 @@ import json -import flask_restful from flask import jsonify, make_response, request from monkey_island.cc.models import Monkey -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services.node import NodeService from monkey_island.cc.services.run_local_monkey import LocalMonkeyRunService -class LocalRun(flask_restful.Resource): +class LocalRun(AbstractResource): + + urls = ["/api/local-monkey"] + @jwt_required def get(self): island_monkey = NodeService.get_monkey_island_monkey() diff --git a/monkey/monkey_island/cc/resources/log.py b/monkey/monkey_island/cc/resources/log.py index 63e4d44f1..541ffb1af 100644 --- a/monkey/monkey_island/cc/resources/log.py +++ b/monkey/monkey_island/cc/resources/log.py @@ -1,17 +1,19 @@ import json -import flask_restful from bson import ObjectId from flask import request from monkey_island.cc.database import mongo -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource from monkey_island.cc.resources.blackbox.utils.telem_store import TestTelemStore +from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services.log import LogService from monkey_island.cc.services.node import NodeService -class Log(flask_restful.Resource): +class Log(AbstractResource): + urls = ["/api/log"] + @jwt_required def get(self): monkey_id = request.args.get("id") diff --git a/monkey/monkey_island/cc/resources/monkey.py b/monkey/monkey_island/cc/resources/monkey.py index c8ff52468..21ffbed36 100644 --- a/monkey/monkey_island/cc/resources/monkey.py +++ b/monkey/monkey_island/cc/resources/monkey.py @@ -1,11 +1,11 @@ import json from datetime import datetime -import flask_restful from flask import request from monkey_island.cc.database import mongo from monkey_island.cc.models.monkey_ttl import create_monkey_ttl_document +from monkey_island.cc.resources.AbstractResource import AbstractResource from monkey_island.cc.resources.blackbox.utils.telem_store import TestTelemStore from monkey_island.cc.resources.utils.semaphores import agent_killing_mutex from monkey_island.cc.server_utils.consts import DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS @@ -16,7 +16,12 @@ from monkey_island.cc.services.node import NodeService # TODO: separate logic from interface -class Monkey(flask_restful.Resource): +class Monkey(AbstractResource): + urls = [ + "/api/agent", + "/api/agent/", + "/api/agent//", + ] # Used by monkey. can't secure. def get(self, guid=None, config_format=None, **kw): diff --git a/monkey/monkey_island/cc/resources/monkey_download.py b/monkey/monkey_island/cc/resources/monkey_download.py index a5750769f..9c19b70dc 100644 --- a/monkey/monkey_island/cc/resources/monkey_download.py +++ b/monkey/monkey_island/cc/resources/monkey_download.py @@ -2,9 +2,9 @@ import hashlib import logging from pathlib import Path -import flask_restful from flask import make_response, send_from_directory +from monkey_island.cc.resources.AbstractResource import AbstractResource from monkey_island.cc.server_utils.consts import MONKEY_ISLAND_ABS_PATH logger = logging.getLogger(__name__) @@ -19,7 +19,8 @@ class UnsupportedOSError(Exception): pass -class MonkeyDownload(flask_restful.Resource): +class MonkeyDownload(AbstractResource): + urls = ["/api/agent/download/"] # Used by monkey. can't secure. def get(self, host_os): diff --git a/monkey/monkey_island/cc/resources/netmap.py b/monkey/monkey_island/cc/resources/netmap.py index a649fff76..c9d99cfb8 100644 --- a/monkey/monkey_island/cc/resources/netmap.py +++ b/monkey/monkey_island/cc/resources/netmap.py @@ -1,11 +1,12 @@ -import flask_restful - -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services.netmap.net_edge import NetEdgeService from monkey_island.cc.services.netmap.net_node import NetNodeService -class NetMap(flask_restful.Resource): +class NetMap(AbstractResource): + urls = ["/api/netmap"] + @jwt_required def get(self, **kw): net_nodes = NetNodeService.get_all_net_nodes() diff --git a/monkey/monkey_island/cc/resources/node.py b/monkey/monkey_island/cc/resources/node.py index d4252354c..e8b6e0079 100644 --- a/monkey/monkey_island/cc/resources/node.py +++ b/monkey/monkey_island/cc/resources/node.py @@ -1,11 +1,13 @@ -import flask_restful from flask import request -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services.node import NodeService -class Node(flask_restful.Resource): +class Node(AbstractResource): + urls = ["/api/netmap/node"] + @jwt_required def get(self): node_id = request.args.get("id") diff --git a/monkey/monkey_island/cc/resources/node_states.py b/monkey/monkey_island/cc/resources/node_states.py index 073aafffd..88b4d7d64 100644 --- a/monkey/monkey_island/cc/resources/node_states.py +++ b/monkey/monkey_island/cc/resources/node_states.py @@ -1,10 +1,11 @@ -import flask_restful - -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services.utils.node_states import NodeStates as NodeStateList -class NodeStates(flask_restful.Resource): +class NodeStates(AbstractResource): + urls = ["/api/netmap/node-states"] + @jwt_required def get(self): return {"node_states": [state.value for state in NodeStateList]} diff --git a/monkey/monkey_island/cc/resources/pba_file_download.py b/monkey/monkey_island/cc/resources/pba_file_download.py index a11e964b1..73658d0ad 100644 --- a/monkey/monkey_island/cc/resources/pba_file_download.py +++ b/monkey/monkey_island/cc/resources/pba_file_download.py @@ -1,14 +1,15 @@ import logging -import flask_restful from flask import make_response, send_file +from monkey_island.cc.resources.AbstractResource import AbstractResource from monkey_island.cc.services import FileRetrievalError, IFileStorageService logger = logging.getLogger(__file__) -class PBAFileDownload(flask_restful.Resource): +class PBAFileDownload(AbstractResource): + urls = ["/api/pba/download/"] """ File download endpoint used by monkey to download user's PBA file """ diff --git a/monkey/monkey_island/cc/resources/pba_file_upload.py b/monkey/monkey_island/cc/resources/pba_file_upload.py index 233cf9e67..9f741a039 100644 --- a/monkey/monkey_island/cc/resources/pba_file_upload.py +++ b/monkey/monkey_island/cc/resources/pba_file_upload.py @@ -1,28 +1,33 @@ import logging from http import HTTPStatus -import flask_restful from flask import Response, make_response, request, send_file from werkzeug.utils import secure_filename as sanitize_filename from common.config_value_paths import PBA_LINUX_FILENAME_PATH, PBA_WINDOWS_FILENAME_PATH -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services import FileRetrievalError, IFileStorageService from monkey_island.cc.services.config import ConfigService logger = logging.getLogger(__file__) - # Front end uses these strings to identify which files to work with (linux or windows) LINUX_PBA_TYPE = "PBAlinux" WINDOWS_PBA_TYPE = "PBAwindows" -class FileUpload(flask_restful.Resource): +class FileUpload(AbstractResource): """ File upload endpoint used to send/receive Custom PBA files """ + urls = [ + "/api/file-upload/", + "/api/file-upload/?load=", + "/api/file-upload/?restore=", + ] + def __init__(self, file_storage_service: IFileStorageService): self._file_storage_service = file_storage_service diff --git a/monkey/monkey_island/cc/resources/propagation_credentials.py b/monkey/monkey_island/cc/resources/propagation_credentials.py index 532501658..340f473b3 100644 --- a/monkey/monkey_island/cc/resources/propagation_credentials.py +++ b/monkey/monkey_island/cc/resources/propagation_credentials.py @@ -1,10 +1,11 @@ -import flask_restful - from monkey_island.cc.database import mongo +from monkey_island.cc.resources.AbstractResource import AbstractResource from monkey_island.cc.services.config import ConfigService -class PropagationCredentials(flask_restful.Resource): +class PropagationCredentials(AbstractResource): + urls = ["/api/propagation-credentials/"] + def get(self, guid: str): monkey_json = mongo.db.monkey.find_one_or_404({"guid": guid}) ConfigService.decrypt_flat_config(monkey_json["config"]) diff --git a/monkey/monkey_island/cc/resources/ransomware_report.py b/monkey/monkey_island/cc/resources/ransomware_report.py index af86e75a1..d229361ea 100644 --- a/monkey/monkey_island/cc/resources/ransomware_report.py +++ b/monkey/monkey_island/cc/resources/ransomware_report.py @@ -1,11 +1,13 @@ -import flask_restful from flask import jsonify -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services.ransomware import ransomware_report -class RansomwareReport(flask_restful.Resource): +class RansomwareReport(AbstractResource): + urls = ["/api/report/ransomware"] + @jwt_required def get(self): return jsonify( diff --git a/monkey/monkey_island/cc/resources/remote_run.py b/monkey/monkey_island/cc/resources/remote_run.py index 8bb0752aa..820419487 100644 --- a/monkey/monkey_island/cc/resources/remote_run.py +++ b/monkey/monkey_island/cc/resources/remote_run.py @@ -1,11 +1,11 @@ import json from typing import Sequence -import flask_restful from botocore.exceptions import ClientError, NoCredentialsError from flask import jsonify, make_response, request -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services import AWSService from monkey_island.cc.services.aws import AWSCommandResults @@ -19,7 +19,9 @@ NO_CREDS_ERROR_FORMAT = ( ) -class RemoteRun(flask_restful.Resource): +class RemoteRun(AbstractResource): + urls = ["/api/remote-monkey"] + def __init__(self, aws_service: AWSService): self._aws_service = aws_service diff --git a/monkey/monkey_island/cc/resources/request_authentication.py b/monkey/monkey_island/cc/resources/request_authentication.py new file mode 100644 index 000000000..6fd97d016 --- /dev/null +++ b/monkey/monkey_island/cc/resources/request_authentication.py @@ -0,0 +1,31 @@ +import logging +from functools import wraps + +import flask_jwt_extended +from flask import make_response +from flask_jwt_extended.exceptions import JWTExtendedException +from jwt import PyJWTError + +logger = logging.getLogger(__name__) + + +def create_access_token(username): + access_token = flask_jwt_extended.create_access_token(identity=username) + logger.debug(f"Created access token for user {username} that begins with {access_token[:4]}") + + return access_token + + +# See https://flask-jwt-extended.readthedocs.io/en/stable/custom_decorators/ +def jwt_required(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + try: + flask_jwt_extended.verify_jwt_in_request() + return fn(*args, **kwargs) + # Catch authentication related errors in the verification or inside the called function. + # All other exceptions propagate + except (JWTExtendedException, PyJWTError) as e: + return make_response({"error": f"Authentication error: {str(e)}"}, 401) + + return wrapper diff --git a/monkey/monkey_island/cc/resources/root.py b/monkey/monkey_island/cc/resources/root.py index aa2087913..23e39aa49 100644 --- a/monkey/monkey_island/cc/resources/root.py +++ b/monkey/monkey_island/cc/resources/root.py @@ -1,10 +1,10 @@ import logging -import flask_restful 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.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services.database import Database from monkey_island.cc.services.infection_lifecycle import get_completed_steps from monkey_island.cc.services.utils.network_utils import local_ip_addresses @@ -12,7 +12,10 @@ from monkey_island.cc.services.utils.network_utils import local_ip_addresses logger = logging.getLogger(__name__) -class Root(flask_restful.Resource): +class Root(AbstractResource): + + urls = ["/api"] + def get(self, action=None): if not action: action = request.args.get("action") diff --git a/monkey/monkey_island/cc/resources/security_report.py b/monkey/monkey_island/cc/resources/security_report.py index b2ce0704e..62b254931 100644 --- a/monkey/monkey_island/cc/resources/security_report.py +++ b/monkey/monkey_island/cc/resources/security_report.py @@ -1,10 +1,11 @@ -import flask_restful - -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services.reporting.report import ReportService -class SecurityReport(flask_restful.Resource): +class SecurityReport(AbstractResource): + urls = ["/api/report/security"] + @jwt_required def get(self): return ReportService.get_report() diff --git a/monkey/monkey_island/cc/resources/telemetry.py b/monkey/monkey_island/cc/resources/telemetry.py index 3358788f3..11e4ba06e 100644 --- a/monkey/monkey_island/cc/resources/telemetry.py +++ b/monkey/monkey_island/cc/resources/telemetry.py @@ -3,21 +3,23 @@ import logging from datetime import datetime import dateutil -import flask_restful from flask import request from monkey_island.cc.database import mongo from monkey_island.cc.models.monkey import Monkey from monkey_island.cc.models.telemetries import get_telemetry_by_query -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource from monkey_island.cc.resources.blackbox.utils.telem_store import TestTelemStore +from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services.node import NodeService from monkey_island.cc.services.telemetry.processing.processing import process_telemetry logger = logging.getLogger(__name__) -class Telemetry(flask_restful.Resource): +class Telemetry(AbstractResource): + urls = ["/api/telemetry", "/api/telemetry/"] + @jwt_required def get(self, **kw): monkey_guid = request.args.get("monkey_guid") diff --git a/monkey/monkey_island/cc/resources/telemetry_feed.py b/monkey/monkey_island/cc/resources/telemetry_feed.py index d880d964d..fa97f3257 100644 --- a/monkey/monkey_island/cc/resources/telemetry_feed.py +++ b/monkey/monkey_island/cc/resources/telemetry_feed.py @@ -3,18 +3,20 @@ from datetime import datetime import dateutil import flask_pymongo -import flask_restful from flask import request from common.common_consts.telem_categories import TelemCategoryEnum from monkey_island.cc.database import mongo -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services.node import NodeService logger = logging.getLogger(__name__) -class TelemetryFeed(flask_restful.Resource): +class TelemetryFeed(AbstractResource): + urls = ["/api/telemetry-feed"] + @jwt_required def get(self, **kw): timestamp = request.args.get("timestamp") diff --git a/monkey/monkey_island/cc/resources/version_update.py b/monkey/monkey_island/cc/resources/version_update.py index 9346bfce4..0544cea4c 100644 --- a/monkey/monkey_island/cc/resources/version_update.py +++ b/monkey/monkey_island/cc/resources/version_update.py @@ -1,14 +1,15 @@ import logging -import flask_restful - from common.version import get_version +from monkey_island.cc.resources.AbstractResource import AbstractResource from monkey_island.cc.services.version_update import VersionUpdateService logger = logging.getLogger(__name__) -class VersionUpdate(flask_restful.Resource): +class VersionUpdate(AbstractResource): + urls = ["/api/version-update"] + def __init__(self): super(VersionUpdate, self).__init__() diff --git a/monkey/monkey_island/cc/resources/zero_trust/finding_event.py b/monkey/monkey_island/cc/resources/zero_trust/finding_event.py index ce99390da..d96e741f9 100644 --- a/monkey/monkey_island/cc/resources/zero_trust/finding_event.py +++ b/monkey/monkey_island/cc/resources/zero_trust/finding_event.py @@ -1,14 +1,15 @@ import json -import flask_restful - -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services.zero_trust.monkey_findings.monkey_zt_finding_service import ( MonkeyZTFindingService, ) -class ZeroTrustFindingEvent(flask_restful.Resource): +class ZeroTrustFindingEvent(AbstractResource): + urls = ["/api/zero-trust/finding-event/"] + @jwt_required def get(self, finding_id: str): return { diff --git a/monkey/monkey_island/cc/resources/zero_trust/zero_trust_report.py b/monkey/monkey_island/cc/resources/zero_trust/zero_trust_report.py index 491b109dc..359328114 100644 --- a/monkey/monkey_island/cc/resources/zero_trust/zero_trust_report.py +++ b/monkey/monkey_island/cc/resources/zero_trust/zero_trust_report.py @@ -3,7 +3,8 @@ import http.client import flask_restful from flask import jsonify -from monkey_island.cc.resources.auth.auth import jwt_required +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.resources.request_authentication import jwt_required from monkey_island.cc.services.zero_trust.zero_trust_report.finding_service import FindingService from monkey_island.cc.services.zero_trust.zero_trust_report.pillar_service import PillarService from monkey_island.cc.services.zero_trust.zero_trust_report.principle_service import ( @@ -15,7 +16,9 @@ REPORT_DATA_FINDINGS = "findings" REPORT_DATA_PRINCIPLES_STATUS = "principles" -class ZeroTrustReport(flask_restful.Resource): +class ZeroTrustReport(AbstractResource): + urls = ["/api/report/zero-trust/"] + @jwt_required def get(self, report_data=None): if report_data == REPORT_DATA_PILLARS: diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/auth/test_auth.py b/monkey/tests/unit_tests/monkey_island/cc/resources/auth/test_auth.py index 8bcc80690..73466a47d 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/auth/test_auth.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/auth/test_auth.py @@ -4,6 +4,7 @@ from unittest.mock import MagicMock import pytest from common.utils.exceptions import IncorrectCredentialsError +from monkey_island.cc.resources.auth.auth import Authenticate USERNAME = "test_user" PASSWORD = "test_password" @@ -22,7 +23,7 @@ def mock_authentication_service(monkeypatch): @pytest.fixture def make_auth_request(flask_client): - url = "/api/auth" + url = Authenticate.urls[0] def inner(request_body): return flask_client.post(url, data=request_body, follow_redirects=True) diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/auth/test_registration.py b/monkey/tests/unit_tests/monkey_island/cc/resources/auth/test_registration.py index 041eec264..e7bab544c 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/auth/test_registration.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/auth/test_registration.py @@ -4,8 +4,9 @@ from unittest.mock import MagicMock import pytest from common.utils.exceptions import AlreadyRegisteredError, InvalidRegistrationCredentialsError +from monkey_island.cc.resources.auth.registration import Registration -REGISTRATION_URL = "/api/registration" +REGISTRATION_URL = Registration.urls[0] USERNAME = "test_user" PASSWORD = "test_password" diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py index 723c777b8..a40766d5e 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py @@ -1,14 +1,12 @@ from unittest.mock import MagicMock import flask_jwt_extended -import flask_restful import pytest -from flask import Flask +from tests.unit_tests.monkey_island.conftest import init_mock_app import monkey_island.cc.app import monkey_island.cc.resources.auth.auth import monkey_island.cc.resources.island_mode -from monkey_island.cc.services.representations import output_json @pytest.fixture @@ -18,7 +16,7 @@ def flask_client(monkeypatch_session): container = MagicMock() container.resolve_dependencies.return_value = [] - with mock_init_app(container).test_client() as client: + with get_mock_app(container).test_client() as client: yield client @@ -27,19 +25,13 @@ def build_flask_client(monkeypatch_session): def inner(container): monkeypatch_session.setattr(flask_jwt_extended, "verify_jwt_in_request", lambda: None) - return mock_init_app(container).test_client() + return get_mock_app(container).test_client() return inner -def mock_init_app(container): - app = Flask(__name__) - app.config["SECRET_KEY"] = "test_key" - - api = flask_restful.Api(app) - api.representations = {"application/json": output_json} - - monkey_island.cc.app.init_app_url_rules(app) +def get_mock_app(container): + app, api = init_mock_app() flask_resource_manager = monkey_island.cc.app.FlaskDIWrapper(api, container) monkey_island.cc.app.init_api_resources(flask_resource_manager) diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py index 37b09aaed..a167468e9 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py @@ -5,6 +5,7 @@ from tests.utils import raise_ from monkey_island.cc.models.island_mode_model import IslandMode from monkey_island.cc.resources import island_mode as island_mode_resource +from monkey_island.cc.resources.island_mode import IslandMode as IslandModeResource @pytest.fixture(scope="function") @@ -19,21 +20,21 @@ def test_island_mode_post(flask_client, mode, monkeypatch): lambda _: None, ) resp = flask_client.post( - "/api/island-mode", data=json.dumps({"mode": mode}), follow_redirects=True + IslandModeResource.urls[0], data=json.dumps({"mode": mode}), follow_redirects=True ) assert resp.status_code == 200 def test_island_mode_post__invalid_mode(flask_client): resp = flask_client.post( - "/api/island-mode", data=json.dumps({"mode": "bogus mode"}), follow_redirects=True + IslandModeResource.urls[0], data=json.dumps({"mode": "bogus mode"}), follow_redirects=True ) assert resp.status_code == 422 @pytest.mark.parametrize("invalid_json", ["42", "{test"]) def test_island_mode_post__invalid_json(flask_client, invalid_json): - resp = flask_client.post("/api/island-mode", data="{test", follow_redirects=True) + resp = flask_client.post(IslandModeResource.urls[0], data="{test", follow_redirects=True) assert resp.status_code == 400 @@ -41,23 +42,25 @@ def test_island_mode_post__internal_server_error(monkeypatch, flask_client): monkeypatch.setattr(island_mode_resource, "set_mode", lambda x: raise_(Exception())) resp = flask_client.post( - "/api/island-mode", data=json.dumps({"mode": "ransomware"}), follow_redirects=True + IslandModeResource.urls[0], data=json.dumps({"mode": "ransomware"}), follow_redirects=True ) assert resp.status_code == 500 @pytest.mark.parametrize("mode", ["ransomware", "advanced"]) def test_island_mode_endpoint(flask_client, uses_database, mode): - flask_client.post("/api/island-mode", data=json.dumps({"mode": mode}), follow_redirects=True) - resp = flask_client.get("/api/island-mode", follow_redirects=True) + flask_client.post( + IslandModeResource.urls[0], data=json.dumps({"mode": mode}), follow_redirects=True + ) + resp = flask_client.get(IslandModeResource.urls[0], follow_redirects=True) assert resp.status_code == 200 assert json.loads(resp.data)["mode"] == mode def test_island_mode_endpoint__invalid_mode(flask_client, uses_database): resp_post = flask_client.post( - "/api/island-mode", data=json.dumps({"mode": "bogus_mode"}), follow_redirects=True + IslandModeResource.urls[0], data=json.dumps({"mode": "bogus_mode"}), follow_redirects=True ) - resp_get = flask_client.get("/api/island-mode", follow_redirects=True) + resp_get = flask_client.get(IslandModeResource.urls[0], follow_redirects=True) assert resp_post.status_code == 422 assert json.loads(resp_get.data)["mode"] is None diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py index 570d3239c..d54d14f84 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py @@ -3,7 +3,9 @@ from typing import BinaryIO import pytest from tests.common import StubDIContainer +from tests.unit_tests.monkey_island.conftest import get_url_for_resource +from monkey_island.cc.resources.pba_file_download import PBAFileDownload from monkey_island.cc.services import FileRetrievalError, IFileStorageService FILE_NAME = "test_file" @@ -40,7 +42,8 @@ def flask_client(build_flask_client): def test_file_download_endpoint(tmp_path, flask_client): - resp = flask_client.get(f"/api/pba/download/{FILE_NAME}") + download_url = get_url_for_resource(PBAFileDownload, filename=FILE_NAME) + resp = flask_client.get(download_url) assert resp.status_code == 200 assert next(resp.response) == FILE_CONTENTS @@ -48,7 +51,8 @@ def test_file_download_endpoint(tmp_path, flask_client): def test_file_download_endpoint_404(tmp_path, flask_client): nonexistant_file_name = "nonexistant_file" + download_url = get_url_for_resource(PBAFileDownload, filename=nonexistant_file_name) - resp = flask_client.get(f"/api/pba/download/{nonexistant_file_name}") + resp = flask_client.get(download_url) assert resp.status_code == 404 diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py index 9cbaa50d8..3ec9a631a 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py @@ -3,9 +3,10 @@ from typing import BinaryIO import pytest from tests.common import StubDIContainer +from tests.unit_tests.monkey_island.conftest import get_url_for_resource from tests.utils import raise_ -from monkey_island.cc.resources.pba_file_upload import LINUX_PBA_TYPE, WINDOWS_PBA_TYPE +from monkey_island.cc.resources.pba_file_upload import LINUX_PBA_TYPE, WINDOWS_PBA_TYPE, FileUpload from monkey_island.cc.services import FileRetrievalError, IFileStorageService TEST_FILE_CONTENTS = b"m0nk3y" @@ -74,8 +75,9 @@ def flask_client(build_flask_client, file_storage_service): @pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) def test_pba_file_upload_post(flask_client, pba_os, mock_set_config_value): + url = get_url_for_resource(FileUpload, target_os=pba_os) resp = flask_client.post( - f"/api/file-upload/{pba_os}", + url, data=TEST_FILE, content_type="multipart/form-data; " "boundary=---------------------------" "1", follow_redirects=True, @@ -84,8 +86,9 @@ def test_pba_file_upload_post(flask_client, pba_os, mock_set_config_value): def test_pba_file_upload_post__invalid(flask_client, mock_set_config_value): + url = get_url_for_resource(FileUpload, target_os="bogus") resp = flask_client.post( - "/api/file-upload/bogus", + url, data=TEST_FILE, content_type="multipart/form-data; " "boundary=---------------------------" "1", follow_redirects=True, @@ -98,9 +101,10 @@ def test_pba_file_upload_post__internal_server_error( flask_client, pba_os, mock_set_config_value, file_storage_service ): file_storage_service.save_file = lambda x, y: raise_(Exception()) + url = get_url_for_resource(FileUpload, target_os=pba_os) resp = flask_client.post( - f"/api/file-upload/{pba_os}", + url, data=TEST_FILE, content_type="multipart/form-data; boundary=---------------------------1", follow_redirects=True, @@ -110,7 +114,8 @@ def test_pba_file_upload_post__internal_server_error( @pytest.mark.parametrize("pba_os", [LINUX_PBA_TYPE, WINDOWS_PBA_TYPE]) def test_pba_file_upload_get__file_not_found(flask_client, pba_os, mock_get_config_value): - resp = flask_client.get(f"/api/file-upload/{pba_os}?load=bogus_mogus.py") + url = get_url_for_resource(FileUpload, target_os=pba_os, filename="bobug_mogus.py") + resp = flask_client.get(url) assert resp.status_code == 404 @@ -118,23 +123,24 @@ def test_pba_file_upload_get__file_not_found(flask_client, pba_os, mock_get_conf def test_pba_file_upload_endpoint( flask_client, pba_os, mock_get_config_value, mock_set_config_value ): + + url_with_os = get_url_for_resource(FileUpload, target_os=pba_os) resp_post = flask_client.post( - f"/api/file-upload/{pba_os}", + url_with_os, data=TEST_FILE, content_type="multipart/form-data; " "boundary=---------------------------" "1", follow_redirects=True, ) - resp_get = flask_client.get(f"/api/file-upload/{pba_os}?load=test.py") + url_with_filename = get_url_for_resource(FileUpload, target_os=pba_os, filename="test.py") + resp_get = flask_client.get(url_with_filename) assert resp_get.status_code == 200 assert resp_get.data == TEST_FILE_CONTENTS # Closing the response closes the file handle, else it can't be deleted resp_get.close() - resp_delete = flask_client.delete( - f"/api/file-upload/{pba_os}", data="test.py", content_type="text/plain;" - ) - resp_get_del = flask_client.get(f"/api/file-upload/{pba_os}?load=test.py") + resp_delete = flask_client.delete(url_with_os, data="test.py", content_type="text/plain;") + resp_get_del = flask_client.get(url_with_filename) assert resp_post.status_code == 200 assert resp_delete.status_code == 200 @@ -145,16 +151,20 @@ def test_pba_file_upload_endpoint( def test_pba_file_upload_endpoint__invalid( flask_client, mock_set_config_value, mock_get_config_value ): + + url_with_os = get_url_for_resource(FileUpload, target_os="bogus") resp_post = flask_client.post( - "/api/file-upload/bogus", + url_with_os, data=TEST_FILE, content_type="multipart/form-data; " "boundary=---------------------------" "1", follow_redirects=True, ) - resp_get = flask_client.get("/api/file-upload/bogus?load=test.py") - resp_delete = flask_client.delete( - "/api/file-upload/bogus", data="test.py", content_type="text/plain;" + + url_with_filename = get_url_for_resource( + FileUpload, target_os="bogus", filename="bobug_mogus.py" ) + resp_get = flask_client.get(url_with_filename) + resp_delete = flask_client.delete(url_with_os, data="test.py", content_type="text/plain;") assert resp_post.status_code == 422 assert resp_get.status_code == 422 assert resp_delete.status_code == 422 diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_remote_run.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_remote_run.py index 17064d355..10ef5fcba 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_remote_run.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_remote_run.py @@ -4,6 +4,7 @@ from unittest.mock import MagicMock import pytest from tests.common import StubDIContainer +from monkey_island.cc.resources import RemoteRun from monkey_island.cc.services import AWSService from monkey_island.cc.services.aws import AWSCommandResults, AWSCommandStatus @@ -23,18 +24,18 @@ def flask_client(build_flask_client, mock_aws_service): def test_get_invalid_action(flask_client): - response = flask_client.get("/api/remote-monkey?action=INVALID") + response = flask_client.get(f"{RemoteRun.urls[0]}?action=INVALID") assert response.text.rstrip() == "{}" def test_get_no_action(flask_client): - response = flask_client.get("/api/remote-monkey") + response = flask_client.get(RemoteRun.urls[0]) assert response.text.rstrip() == "{}" def test_get_not_aws(flask_client, mock_aws_service): mock_aws_service.island_is_running_on_aws = MagicMock(return_value=False) - response = flask_client.get("/api/remote-monkey?action=list_aws") + response = flask_client.get(f"{RemoteRun.urls[0]}?action=list_aws") assert response.text.rstrip() == '{"is_aws":false}' @@ -47,7 +48,7 @@ def test_get_instances(flask_client, mock_aws_service): mock_aws_service.island_is_running_on_aws = MagicMock(return_value=True) mock_aws_service.get_managed_instances = MagicMock(return_value=instances) - response = flask_client.get("/api/remote-monkey?action=list_aws") + response = flask_client.get(f"{RemoteRun.urls[0]}?action=list_aws") assert json.loads(response.text)["instances"] == instances assert json.loads(response.text)["is_aws"] is True @@ -57,12 +58,12 @@ def test_get_instances(flask_client, mock_aws_service): def test_post_no_type(flask_client): - response = flask_client.post("/api/remote-monkey", data="{}") + response = flask_client.post(RemoteRun.urls[0], data="{}") assert response.status_code == 500 def test_post_invalid_type(flask_client): - response = flask_client.post("/api/remote-monkey", data='{"type": "INVALID"}') + response = flask_client.post(RemoteRun.urls[0], data='{"type": "INVALID"}') assert response.status_code == 500 @@ -103,6 +104,6 @@ def test_post(flask_client, mock_aws_service): }, ] - response = flask_client.post("/api/remote-monkey", data=request_body) + response = flask_client.post(RemoteRun.urls[0], data=request_body) assert json.loads(response.text)["result"] == expected_result diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py b/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py index aa4cfdb4b..83c9d2dde 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/aws/test_aws_command_runner.py @@ -14,15 +14,6 @@ from monkey_island.cc.services.aws.aws_command_runner import ( TIMEOUT = 0.03 INSTANCE_ID = "BEEFFACE" ISLAND_IP = "127.0.0.1" -""" - "commands": [ - "wget --no-check-certificate " - "https://172.31.32.78:5000/api/agent/download/linux " - "-O monkey-linux-64; chmod +x " - "monkey-linux-64; ./monkey-linux-64 " - "m0nk3y -s 172.31.32.78:5000" - ] - """ @pytest.fixture diff --git a/monkey/tests/unit_tests/monkey_island/cc/test_app.py b/monkey/tests/unit_tests/monkey_island/cc/test_app.py new file mode 100644 index 000000000..225f87ff8 --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/test_app.py @@ -0,0 +1,83 @@ +import pytest +from tests.common import StubDIContainer +from tests.unit_tests.monkey_island.conftest import mock_flask_resource_manager + +from monkey_island.cc.app import FlaskDIWrapper +from monkey_island.cc.resources.AbstractResource import AbstractResource + + +def get_mock_resource(name, urls): + class MockResource(AbstractResource): + urls = [] + + def get(self, something=None): + pass + + mock = type(name, MockResource.__bases__, dict(MockResource.__dict__)) + mock.urls = urls + return mock + + +@pytest.fixture +def resource_manager(): + container = StubDIContainer() + return mock_flask_resource_manager(container) + + +def test_duplicate_urls(resource_manager): + resource = get_mock_resource("res1", ["/url"]) + + resource2 = get_mock_resource("res1", ["/new_url", "/url"]) + + resource_manager.add_resource(resource) + with pytest.raises(FlaskDIWrapper.DuplicateURLError): + resource_manager.add_resource(resource2) + + +def test_duplicate_urls__parameters(resource_manager): + resource1 = get_mock_resource("res1", ["/url/"]) + resource2 = get_mock_resource("res2", ["/url/"]) + + resource_manager.add_resource(resource1) + with pytest.raises(FlaskDIWrapper.DuplicateURLError): + resource_manager.add_resource(resource2) + + +def test_duplicate_urls__multiple_parameters(resource_manager): + resource1 = get_mock_resource("res1", ["/url//"]) + resource2 = get_mock_resource("res2", ["/url//"]) + + resource_manager.add_resource(resource1) + with pytest.raises(FlaskDIWrapper.DuplicateURLError): + resource_manager.add_resource(resource2) + + +def test_adding_resources(resource_manager): + resource = get_mock_resource("res1", ["/url"]) + + resource2 = get_mock_resource("res2", ["/different_url", "/another_different"]) + + resource3 = get_mock_resource("res3", ["/yet_another/"]) + + # Following shouldn't raise an exception + resource_manager.add_resource(resource) + resource_manager.add_resource(resource2) + resource_manager.add_resource(resource3) + + +def test_url_check_slash_stripping__trailing_slash(resource_manager): + resource = get_mock_resource("res", ["/url"]) + resource2 = get_mock_resource("res2", ["/url/"]) + + resource_manager.add_resource(resource) + with pytest.raises(FlaskDIWrapper.DuplicateURLError): + resource_manager.add_resource(resource2) + + +def test_url_check_slash_stripping__path_separation(resource_manager): + resource3 = get_mock_resource("res3", ["/beef/face/"]) + resource4 = get_mock_resource("res4", ["/beefface"]) + + # Following shouldn't raise and exception + resource_manager.add_resource(resource3) + resource_manager.add_resource(resource4) diff --git a/monkey/tests/unit_tests/monkey_island/conftest.py b/monkey/tests/unit_tests/monkey_island/conftest.py index 2ccecd616..094bbb2aa 100644 --- a/monkey/tests/unit_tests/monkey_island/conftest.py +++ b/monkey/tests/unit_tests/monkey_island/conftest.py @@ -1,7 +1,15 @@ import os +import re from collections.abc import Callable +from typing import Set +import flask_restful import pytest +from flask import Flask + +import monkey_island +from monkey_island.cc.resources.AbstractResource import AbstractResource +from monkey_island.cc.services.representations import output_json @pytest.fixture(scope="module") @@ -19,3 +27,45 @@ def create_empty_tmp_file(tmpdir: str) -> Callable: return new_file return inner + + +def init_mock_app(): + app = Flask(__name__) + app.config["SECRET_KEY"] = "test_key" + + api = flask_restful.Api(app) + api.representations = {"application/json": output_json} + + monkey_island.cc.app.init_app_url_rules(app) + return app, api + + +def mock_flask_resource_manager(container): + _, api = init_mock_app() + flask_resource_manager = monkey_island.cc.app.FlaskDIWrapper(api, container) + + return flask_resource_manager + + +def get_url_for_resource(resource: AbstractResource, **kwargs): + chosen_url = None + for url in resource.urls: + if _get_url_keywords(url) == set(kwargs.keys()): + chosen_url = url + if not chosen_url: + raise Exception( + f"Resource {resource} doesn't contain a url that matches {kwargs} keywords." + ) + + for key, value in kwargs.items(): + reg_pattern = f"<.*:{key}>" + chosen_url = re.sub(pattern=reg_pattern, repl=value, string=chosen_url) + + return chosen_url + + +def _get_url_keywords(url: str) -> Set[str]: + # Match pattern , but only put "keyword" in a group + reg_pattern = "(?:<.*?:)(.*?)(?:>)" + reg_matches = re.finditer(reg_pattern, url) + return set([match.groups()[0] for match in reg_matches])