diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index 07e08ea9a..be135b494 100644 --- a/monkey/monkey_island/cc/app.py +++ b/monkey/monkey_island/cc/app.py @@ -1,11 +1,12 @@ import os import uuid from datetime import timedelta -from typing import Type +from typing import Sequence, Type import flask_restful from flask import Flask, Response, send_from_directory from werkzeug.exceptions import NotFound +from werkzeug.routing import Map from common import DIContainer from monkey_island.cc.database import database, mongo @@ -25,6 +26,7 @@ from monkey_island.cc.resources.configuration_import import ConfigurationImport from monkey_island.cc.resources.edge import Edge from monkey_island.cc.resources.exploitations.manual_exploitation import ManualExploitation from monkey_island.cc.resources.exploitations.monkey_exploitation import MonkeyExploitation +from monkey_island.cc.resources.i_resource import IResource from monkey_island.cc.resources.island_configuration import IslandConfiguration from monkey_island.cc.resources.island_logs import IslandLog from monkey_island.cc.resources.island_mode import IslandMode @@ -109,13 +111,35 @@ def init_app_url_rules(app): class FlaskDIWrapper: + class URLAlreadyExistsError(Exception): + pass + 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): + def add_resource(self, resource: Type[IResource]): dependencies = self._container.resolve_dependencies(resource) - self._api.add_resource(resource, *urls, resource_class_args=dependencies) + FlaskDIWrapper._check_for_duplicate_urls(self._api.app.url_map, resource.urls) + self._api.add_resource(resource, *resource.urls, resource_class_args=dependencies) + + @staticmethod + def _check_for_duplicate_urls(url_map: Map, urls: Sequence[str]): + for url in urls: + if FlaskDIWrapper._is_url_added(url_map, url): + raise FlaskDIWrapper.URLAlreadyExistsError( + f"URL {url} has already been registered!" + ) + + @staticmethod + def _is_url_added(url_map: Map, url_to_add: str) -> bool: + return bool( + [ + registered_url + for registered_url in url_map.iter_rules() + if str(registered_url).strip("/") == url_to_add.strip("/") + ] + ) def init_api_resources(api: FlaskDIWrapper): diff --git a/monkey/monkey_island/cc/resources/i_resource.py b/monkey/monkey_island/cc/resources/i_resource.py new file mode 100644 index 000000000..239b37854 --- /dev/null +++ b/monkey/monkey_island/cc/resources/i_resource.py @@ -0,0 +1,17 @@ +from abc import ABCMeta, abstractmethod +from typing import Sequence + +from flask.views import MethodViewType + + +# Flask resources inherit from flask_restful.Resource, so custom interface +# must implement both metaclasses +class AbstractResource(ABCMeta, MethodViewType): + pass + + +class IResource(metaclass=AbstractResource): + @property + @abstractmethod + def urls(self) -> Sequence[str]: + pass 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..61d596472 --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/test_app.py @@ -0,0 +1,62 @@ +import flask_restful +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.i_resource import IResource + + +def get_mock_resource(name, urls): + class MockResource(flask_restful.Resource, IResource): + urls = [] + + def get(self, something=None): + pass + + mock = type(name, MockResource.__bases__, dict(MockResource.__dict__)) + mock.urls = urls + return mock + + +@pytest.fixture +def resource_mng(): + container = StubDIContainer() + return mock_flask_resource_manager(container) + + +def test_duplicate_urls(resource_mng): + resource = get_mock_resource("res1", ["/url"]) + + resource2 = get_mock_resource("res1", ["/new_url", "/url"]) + + resource_mng.add_resource(resource) + with pytest.raises(FlaskDIWrapper.URLAlreadyExistsError): + resource_mng.add_resource(resource2) + + +def test_adding_resources(resource_mng): + resource = get_mock_resource("res1", ["/url"]) + + resource2 = get_mock_resource("res2", ["/different_url", "/another_different"]) + + resource3 = get_mock_resource("res3", ["/yet_another/"]) + + resource_mng.add_resource(resource) + resource_mng.add_resource(resource2) + resource_mng.add_resource(resource3) + + +def test_url_check_slash_stripping(resource_mng): + resource = get_mock_resource("res", ["/url"]) + resource2 = get_mock_resource("res2", ["/url/"]) + + resource_mng.add_resource(resource) + with pytest.raises(FlaskDIWrapper.URLAlreadyExistsError): + resource_mng.add_resource(resource2) + + resource3 = get_mock_resource("res3", ["/beef/face/"]) + resource4 = get_mock_resource("res4", ["/beefface"]) + + resource_mng.add_resource(resource3) + resource_mng.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..66780882c 100644 --- a/monkey/tests/unit_tests/monkey_island/conftest.py +++ b/monkey/tests/unit_tests/monkey_island/conftest.py @@ -1,7 +1,12 @@ import os from collections.abc import Callable +import flask_restful import pytest +from flask import Flask + +import monkey_island +from monkey_island.cc.services.representations import output_json @pytest.fixture(scope="module") @@ -19,3 +24,16 @@ def create_empty_tmp_file(tmpdir: str) -> Callable: return new_file return inner + + +def mock_flask_resource_manager(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) + flask_resource_manager = monkey_island.cc.app.FlaskDIWrapper(api, container) + + return flask_resource_manager