From 6fa52d0637591cdcc27df73048d369e0c53003ba Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Thu, 30 Jun 2022 12:56:50 -0400
Subject: [PATCH] Island: Remove ModeNotSetError

---
 monkey/monkey_island/cc/resources/island_mode.py    | 10 +++-------
 monkey/monkey_island/cc/services/config.py          | 13 ++++++++-----
 .../cc/services/config_manipulators.py              |  1 +
 .../cc/services/mode/island_mode_service.py         |  9 ++-------
 4 files changed, 14 insertions(+), 19 deletions(-)

diff --git a/monkey/monkey_island/cc/resources/island_mode.py b/monkey/monkey_island/cc/resources/island_mode.py
index 6978f3b14..b4d358b66 100644
--- a/monkey/monkey_island/cc/resources/island_mode.py
+++ b/monkey/monkey_island/cc/resources/island_mode.py
@@ -6,7 +6,7 @@ from flask import make_response, request
 from monkey_island.cc.resources.AbstractResource import AbstractResource
 from monkey_island.cc.resources.request_authentication import jwt_required
 from monkey_island.cc.services.config_manipulator import update_config_on_mode_set
-from monkey_island.cc.services.mode.island_mode_service import ModeNotSetError, get_mode, set_mode
+from monkey_island.cc.services.mode.island_mode_service import get_mode, set_mode
 from monkey_island.cc.services.mode.mode_enum import IslandModeEnum
 
 logger = logging.getLogger(__name__)
@@ -40,9 +40,5 @@ class IslandMode(AbstractResource):
 
     @jwt_required
     def get(self):
-        try:
-            island_mode = get_mode()
-            return make_response({"mode": island_mode}, 200)
-
-        except ModeNotSetError:
-            return make_response({"mode": None}, 200)
+        island_mode = get_mode()
+        return make_response({"mode": island_mode}, 200)
diff --git a/monkey/monkey_island/cc/services/config.py b/monkey/monkey_island/cc/services/config.py
index e0be40117..3cb5ce09f 100644
--- a/monkey/monkey_island/cc/services/config.py
+++ b/monkey/monkey_island/cc/services/config.py
@@ -25,7 +25,8 @@ from monkey_island.cc.server_utils.encryption import (
 )
 from monkey_island.cc.services.config_manipulator import update_config_per_mode
 from monkey_island.cc.services.config_schema.config_schema import SCHEMA
-from monkey_island.cc.services.mode.island_mode_service import ModeNotSetError, get_mode
+from monkey_island.cc.services.mode.island_mode_service import get_mode
+from monkey_island.cc.services.mode.mode_enum import IslandModeEnum
 from monkey_island.cc.services.post_breach_files import PostBreachFilesService
 
 logger = logging.getLogger(__name__)
@@ -250,11 +251,13 @@ class ConfigService:
     def reset_config():
         PostBreachFilesService.remove_PBA_files()
         config = ConfigService.get_default_config(True)
-        try:
-            mode = get_mode()
-            update_config_per_mode(mode, config, should_encrypt=False)
-        except ModeNotSetError:
+
+        mode = get_mode()
+        if mode == IslandModeEnum.UNSET.value:
             ConfigService.update_config(config, should_encrypt=False)
+        else:
+            update_config_per_mode(mode, config, should_encrypt=False)
+
         logger.info("Monkey config reset was called")
 
     @staticmethod
diff --git a/monkey/monkey_island/cc/services/config_manipulators.py b/monkey/monkey_island/cc/services/config_manipulators.py
index 291892371..947a32971 100644
--- a/monkey/monkey_island/cc/services/config_manipulators.py
+++ b/monkey/monkey_island/cc/services/config_manipulators.py
@@ -3,4 +3,5 @@ from monkey_island.cc.services.mode.mode_enum import IslandModeEnum
 MANIPULATOR_PER_MODE = {
     IslandModeEnum.ADVANCED.value: {},
     IslandModeEnum.RANSOMWARE.value: {"monkey.post_breach.post_breach_actions": []},
+    IslandModeEnum.UNSET.value: {},
 }
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 b745ebef1..32a3d943d 100644
--- a/monkey/monkey_island/cc/services/mode/island_mode_service.py
+++ b/monkey/monkey_island/cc/services/mode/island_mode_service.py
@@ -3,6 +3,7 @@ from monkey_island.cc.services.mode.mode_enum import IslandModeEnum
 
 
 def set_mode(mode: IslandModeEnum):
+    IslandMode.drop_collection()
     island_mode_model = IslandMode()
     island_mode_model.mode = mode.value
     island_mode_model.save()
@@ -13,10 +14,4 @@ def get_mode() -> str:
         mode = IslandMode.objects[0].mode
         return mode
     else:
-        raise ModeNotSetError
-
-
-class ModeNotSetError(Exception):
-    """
-    Throw this exception when island mode is not set.
-    """
+        return IslandModeEnum.UNSET.value