Merge pull request #1321 from guardicore/1241/ransomware-quickstart-start-over

Ransomware quickstart unset mode to get method in island mode

monkey/monkey_island/cc/resources/island_mode.py

@@ -5,7 +5,7 @@ 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 get_mode, set_mode
+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
 
 logger = logging.getLogger(__name__)
@@ -32,5 +32,6 @@ class IslandMode(flask_restful.Resource):
         try:
             island_mode = get_mode()
             return make_response({"mode": island_mode}, 200)
-        except IndexError:
-            return make_response({}, 422)
+
+        except ModeNotSetError:
+            return make_response({"mode": None}, 200)

monkey/monkey_island/cc/services/mode/island_mode_service.py

@@ -8,6 +8,15 @@ def set_mode(mode: IslandModeEnum):
     island_mode_model.save()
 
 
-def get_mode():
-    mode = IslandMode.objects[0].mode
-    return mode
+def get_mode() -> str:
+    if IslandMode.objects:
+        mode = IslandMode.objects[0].mode
+        return mode
+    else:
+        raise ModeNotSetError
+
+
+class ModeNotSetError(Exception):
+    """
+    Throw this exception when island mode is not set.
+    """

monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py

@@ -51,8 +51,9 @@ def test_island_mode_endpoint(flask_client, uses_database, mode):
 
 
 def test_island_mode_endpoint__invalid_mode(flask_client, uses_database):
-    flask_client.post(
+    resp_post = flask_client.post(
         "/api/island-mode", data=json.dumps({"mode": "bogus_mode"}), follow_redirects=True
     )
-    resp = flask_client.get("/api/island-mode", follow_redirects=True)
-    assert resp.status_code == 422
+    resp_get = flask_client.get("/api/island-mode", follow_redirects=True)
+    assert resp_post.status_code == 422
+    assert json.loads(resp_get.data)["mode"] is None