Island: Define IResource interface and check for duplicate URL's

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

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

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/<string:something>"])
+
+    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)

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