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 index d0a109564..cf8a3679e 100644 --- a/monkey/monkey_island/cc/resources/island_mode.py +++ b/monkey/monkey_island/cc/resources/island_mode.py @@ -1,20 +1,24 @@ 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 import island_mode_service +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): body = json.loads(request.data) mode_str = body.get("mode") - mode = IslandModeEnum(mode_str) - island_mode_service.set_mode(mode) - - # TODO return status - return make_response({}) + try: + mode = IslandModeEnum(mode_str) + mode_value = set_mode(mode) + return make_response({"status": "MODE_FOUND", "mode": mode_value}, 200) + except ValueError: + return make_response({"status": "MODE_NOT_FOUND"}, 404) diff --git a/monkey/monkey_island/cc/services/mode/island_mode_service.py b/monkey/monkey_island/cc/services/mode/island_mode_service.py index 86f3e6d09..75041d2c1 100644 --- a/monkey/monkey_island/cc/services/mode/island_mode_service.py +++ b/monkey/monkey_island/cc/services/mode/island_mode_service.py @@ -1,5 +1,9 @@ +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): - pass + island_mode_model = IslandMode() + island_mode_model.mode = mode.value + island_mode_model.save() + return mode.value 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 0e82fe163..2421e00aa 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py @@ -1,11 +1,13 @@ import flask_jwt_extended import flask_restful +import mongoengine 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.models.island_mode_model import IslandMode from monkey_island.cc.services.representations import output_json @@ -30,3 +32,14 @@ def mock_init_app(): monkey_island.cc.app.init_api_resources(api) return app + + +@pytest.fixture(scope="module", autouse=True) +def fake_mongo(): + mongoengine.disconnect() + mongoengine.connect("mongoenginetest", host="mongomock://localhost") + + +@pytest.fixture(scope="function") +def uses_database(): + IslandMode.objects().delete() 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 01cb17aea..13ded4799 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 @@ -2,12 +2,34 @@ import json import pytest +from monkey_island.cc.models.island_mode_model import IslandMode + def test_island_mode_post(flask_client): - resp = flask_client.post('/api/island-mode', data=json.dumps({"mode": "ransomware"}), follow_redirects=True) + resp = flask_client.post( + "/api/island-mode", data=json.dumps({"mode": "ransomware"}), follow_redirects=True + ) assert resp.status_code == 200 def test_island_mode_post__invalid_mode(flask_client): - with pytest.raises(TypeError): - flask_client.post('/api/island-mode', data=json.dumps({"mode": "bogus mode"}), follow_redirects=True) + resp = flask_client.post( + "/api/island-mode", data=json.dumps({"mode": "bogus mode"}), follow_redirects=True + ) + assert resp.status_code == 404 + + +@pytest.mark.usefixtures("uses_database") +def test_island_mode_post__set_model(flask_client): + flask_client.post( + "/api/island-mode", data=json.dumps({"mode": "ransomware"}), follow_redirects=True + ) + assert IslandMode.objects[0].mode == "ransomware" + + +@pytest.mark.usefixtures("uses_database") +def test_island_mode_post__set_invalid_model(flask_client): + flask_client.post( + "/api/island-mode", data=json.dumps({"mode": "bogus mode"}), follow_redirects=True + ) + assert len(IslandMode.objects) == 0