diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index 8800d382a..817a43333 100644 --- a/monkey/monkey_island/cc/app.py +++ b/monkey/monkey_island/cc/app.py @@ -26,6 +26,7 @@ from monkey_island.cc.resources.edge import Edge from monkey_island.cc.resources.environment import Environment 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 from monkey_island.cc.resources.local_run import LocalRun from monkey_island.cc.resources.log import Log from monkey_island.cc.resources.monkey import Monkey @@ -132,6 +133,8 @@ def init_api_resources(api): api.add_resource( Telemetry, "/api/telemetry", "/api/telemetry/", "/api/telemetry/" ) + + api.add_resource(IslandMode, "/api/island-mode") api.add_resource(MonkeyConfiguration, "/api/configuration", "/api/configuration/") api.add_resource(IslandConfiguration, "/api/configuration/island", "/api/configuration/island/") api.add_resource(ConfigurationExport, "/api/configuration/export") diff --git a/monkey/monkey_island/cc/models/island_mode_model.py b/monkey/monkey_island/cc/models/island_mode_model.py new file mode 100644 index 000000000..dec93e501 --- /dev/null +++ b/monkey/monkey_island/cc/models/island_mode_model.py @@ -0,0 +1,5 @@ +from mongoengine import Document, StringField + + +class IslandMode(Document): + mode = StringField() diff --git a/monkey/monkey_island/cc/resources/island_mode.py b/monkey/monkey_island/cc/resources/island_mode.py new file mode 100644 index 000000000..51df90e2d --- /dev/null +++ b/monkey/monkey_island/cc/resources/island_mode.py @@ -0,0 +1,28 @@ +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.services.mode.island_mode_service import set_mode +from monkey_island.cc.services.mode.mode_enum import IslandModeEnum + +logger = logging.getLogger(__name__) + + +class IslandMode(flask_restful.Resource): + @jwt_required + def post(self): + try: + body = json.loads(request.data) + mode_str = body.get("mode") + + mode = IslandModeEnum(mode_str) + set_mode(mode) + + return make_response({}, 200) + except (AttributeError, json.decoder.JSONDecodeError): + return make_response({}, 400) + except ValueError: + return make_response({}, 422) diff --git a/monkey/monkey_island/cc/services/mode/__init__.py b/monkey/monkey_island/cc/services/mode/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/monkey/monkey_island/cc/services/mode/island_mode_service.py b/monkey/monkey_island/cc/services/mode/island_mode_service.py new file mode 100644 index 000000000..05195e301 --- /dev/null +++ b/monkey/monkey_island/cc/services/mode/island_mode_service.py @@ -0,0 +1,8 @@ +from monkey_island.cc.models.island_mode_model import IslandMode +from monkey_island.cc.services.mode.mode_enum import IslandModeEnum + + +def set_mode(mode: IslandModeEnum): + island_mode_model = IslandMode() + island_mode_model.mode = mode.value + island_mode_model.save() diff --git a/monkey/monkey_island/cc/services/mode/mode_enum.py b/monkey/monkey_island/cc/services/mode/mode_enum.py new file mode 100644 index 000000000..fce46db97 --- /dev/null +++ b/monkey/monkey_island/cc/services/mode/mode_enum.py @@ -0,0 +1,6 @@ +from enum import Enum + + +class IslandModeEnum(Enum): + RANSOMWARE = "ransomware" + ADVANCED = "advanced" diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py new file mode 100644 index 000000000..0e82fe163 --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py @@ -0,0 +1,32 @@ +import flask_jwt_extended +import flask_restful +import pytest +from flask import Flask + +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 + + +# We can't scope to module, because monkeypatch is a function scoped decorator. +# Potential solutions: https://github.com/pytest-dev/pytest/issues/363#issuecomment-406536200 or +# https://stackoverflow.com/questions/53963822/python-monkeypatch-setattr-with-pytest-fixture-at-module-scope +@pytest.fixture(scope="function") +def flask_client(monkeypatch): + monkeypatch.setattr(flask_jwt_extended, "verify_jwt_in_request", lambda: None) + + with mock_init_app().test_client() as client: + yield client + + +def mock_init_app(): + app = Flask(__name__) + + api = flask_restful.Api(app) + api.representations = {"application/json": output_json} + + monkey_island.cc.app.init_app_url_rules(app) + monkey_island.cc.app.init_api_resources(api) + + return app 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 new file mode 100644 index 000000000..b8d6a84c2 --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py @@ -0,0 +1,56 @@ +import json + +import pytest +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 + + +@pytest.fixture(scope="function") +def uses_database(): + IslandMode.objects().delete() + + +@pytest.mark.parametrize("mode", ["ransomware", "advanced"]) +def test_island_mode_post(flask_client, mode): + resp = flask_client.post( + "/api/island-mode", 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 + ) + 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) + assert resp.status_code == 400 + + +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 + ) + assert resp.status_code == 500 + + +def test_island_mode_post__set_model(flask_client, uses_database): + flask_client.post( + "/api/island-mode", data=json.dumps({"mode": "ransomware"}), follow_redirects=True + ) + assert IslandMode.objects[0].mode == "ransomware" + + +def test_island_mode_post__set_invalid_model(flask_client, uses_database): + flask_client.post( + "/api/island-mode", data=json.dumps({"mode": "bogus mode"}), follow_redirects=True + ) + assert len(IslandMode.objects) == 0 diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_post_breach_files.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_post_breach_files.py index 5a2ddaa17..90a649a39 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/test_post_breach_files.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/test_post_breach_files.py @@ -2,15 +2,12 @@ import os import pytest from tests.monkey_island.utils import assert_windows_permissions +from tests.utils import raise_ from monkey_island.cc.server_utils.file_utils import is_windows_os from monkey_island.cc.services.post_breach_files import PostBreachFilesService -def raise_(ex): - raise ex - - @pytest.fixture(autouse=True) def custom_pba_directory(tmpdir): PostBreachFilesService.initialize(tmpdir) diff --git a/monkey/tests/utils.py b/monkey/tests/utils.py index 2be032aad..8aea2d007 100644 --- a/monkey/tests/utils.py +++ b/monkey/tests/utils.py @@ -18,3 +18,7 @@ def hash_file(filepath: Path): sha256.update(block) return sha256.hexdigest() + + +def raise_(ex): + raise ex diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 618fabaa6..5a430dc6c 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -171,6 +171,7 @@ MONKEY_LINUX_RUNNING # unused variable (monkey/monkey_island/cc/services/utils/ import_status # monkey_island\cc\resources\configuration_import.py:19 config_schema # monkey_island\cc\resources\configuration_import.py:25 exception_stream # unused attribute (monkey_island/cc/server_setup.py:104) +ADVANCED # unused attribute (monkey/monkey_island/cc/services/mode/mode_enum.py:6:) # these are not needed for it to work, but may be useful extra information to understand what's going on WINDOWS_PBA_TYPE # unused variable (monkey/monkey_island/cc/resources/pba_file_upload.py:23)