Island: Remove ModeNotSetError

This commit is contained in:
Mike Salvatore 2022-06-30 12:56:50 -04:00
parent d4c7b97229
commit 6fa52d0637
4 changed files with 14 additions and 19 deletions

View File

@ -6,7 +6,7 @@ from flask import make_response, request
from monkey_island.cc.resources.AbstractResource import AbstractResource from monkey_island.cc.resources.AbstractResource import AbstractResource
from monkey_island.cc.resources.request_authentication import jwt_required 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.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 from monkey_island.cc.services.mode.mode_enum import IslandModeEnum
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -40,9 +40,5 @@ class IslandMode(AbstractResource):
@jwt_required @jwt_required
def get(self): def get(self):
try:
island_mode = get_mode() island_mode = get_mode()
return make_response({"mode": island_mode}, 200) return make_response({"mode": island_mode}, 200)
except ModeNotSetError:
return make_response({"mode": None}, 200)

View File

@ -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_manipulator import update_config_per_mode
from monkey_island.cc.services.config_schema.config_schema import SCHEMA 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 from monkey_island.cc.services.post_breach_files import PostBreachFilesService
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -250,11 +251,13 @@ class ConfigService:
def reset_config(): def reset_config():
PostBreachFilesService.remove_PBA_files() PostBreachFilesService.remove_PBA_files()
config = ConfigService.get_default_config(True) config = ConfigService.get_default_config(True)
try:
mode = get_mode() mode = get_mode()
update_config_per_mode(mode, config, should_encrypt=False) if mode == IslandModeEnum.UNSET.value:
except ModeNotSetError:
ConfigService.update_config(config, should_encrypt=False) ConfigService.update_config(config, should_encrypt=False)
else:
update_config_per_mode(mode, config, should_encrypt=False)
logger.info("Monkey config reset was called") logger.info("Monkey config reset was called")
@staticmethod @staticmethod

View File

@ -3,4 +3,5 @@ from monkey_island.cc.services.mode.mode_enum import IslandModeEnum
MANIPULATOR_PER_MODE = { MANIPULATOR_PER_MODE = {
IslandModeEnum.ADVANCED.value: {}, IslandModeEnum.ADVANCED.value: {},
IslandModeEnum.RANSOMWARE.value: {"monkey.post_breach.post_breach_actions": []}, IslandModeEnum.RANSOMWARE.value: {"monkey.post_breach.post_breach_actions": []},
IslandModeEnum.UNSET.value: {},
} }

View File

@ -3,6 +3,7 @@ from monkey_island.cc.services.mode.mode_enum import IslandModeEnum
def set_mode(mode: IslandModeEnum): def set_mode(mode: IslandModeEnum):
IslandMode.drop_collection()
island_mode_model = IslandMode() island_mode_model = IslandMode()
island_mode_model.mode = mode.value island_mode_model.mode = mode.value
island_mode_model.save() island_mode_model.save()
@ -13,10 +14,4 @@ def get_mode() -> str:
mode = IslandMode.objects[0].mode mode = IslandMode.objects[0].mode
return mode return mode
else: else:
raise ModeNotSetError return IslandModeEnum.UNSET.value
class ModeNotSetError(Exception):
"""
Throw this exception when island mode is not set.
"""