From 2a1d41f6c7f1070caa07595bf1eff8dde1921b94 Mon Sep 17 00:00:00 2001
From: Ilija Lazoroski <ilija.la@live.com>
Date: Wed, 14 Jul 2021 12:40:08 +0200
Subject: [PATCH] Island: Add unset mode to get method in island mode

---
 monkey/monkey_island/cc/resources/island_mode.py  |  7 ++++---
 .../cc/services/mode/island_mode_service.py       | 15 ++++++++++++---
 .../cc/resources/test_island_mode.py              |  7 ++++---
 3 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/monkey/monkey_island/cc/resources/island_mode.py b/monkey/monkey_island/cc/resources/island_mode.py
index 77e1e6953..7698fca9d 100644
--- a/monkey/monkey_island/cc/resources/island_mode.py
+++ b/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)
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 e850c64a7..b745ebef1 100644
--- a/monkey/monkey_island/cc/services/mode/island_mode_service.py
+++ b/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.
+    """
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 f4fb8a797..823f9dad5 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
@@ -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