From 0681782505aaa8d5f0997dd3c1205835407ed010 Mon Sep 17 00:00:00 2001 From: Shreya Date: Tue, 20 Jul 2021 13:16:06 +0530 Subject: [PATCH] island: Pass `mode` as an enum instead of enum value when updating config, and update typehints --- monkey/monkey_island/cc/services/config.py | 9 +++++---- .../cc/services/mode/set_island_mode_service.py | 2 +- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/monkey/monkey_island/cc/services/config.py b/monkey/monkey_island/cc/services/config.py index 82ab04864..e9c3598a3 100644 --- a/monkey/monkey_island/cc/services/config.py +++ b/monkey/monkey_island/cc/services/config.py @@ -13,6 +13,7 @@ from monkey_island.cc.server_utils.encryptor import get_encryptor from monkey_island.cc.services.config_filters import FILTER_PER_MODE from monkey_island.cc.services.config_schema.config_schema import SCHEMA from monkey_island.cc.services.mode.get_island_mode_service import ModeNotSetError, 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.utils.network_utils import local_ip_addresses @@ -248,18 +249,18 @@ class ConfigService: return config @staticmethod - def update_config_on_mode_set(mode: str) -> None: + def update_config_on_mode_set(mode: IslandModeEnum) -> None: config = ConfigService.get_config() ConfigService.update_config_per_mode(mode, config, True) @staticmethod - def update_config_per_mode(mode: str, config: Dict, should_encrypt: bool) -> None: + def update_config_per_mode(mode: IslandModeEnum, config: Dict, should_encrypt: bool) -> None: config = ConfigService._set_default_config_values_per_mode(mode, config) ConfigService.update_config(config_json=config, should_encrypt=True) @staticmethod - def _set_default_config_values_per_mode(mode: str, config: Dict) -> Dict: - config_filter = FILTER_PER_MODE[mode] + def _set_default_config_values_per_mode(mode: IslandModeEnum, config: Dict) -> Dict: + config_filter = FILTER_PER_MODE[mode.value] config = ConfigService._apply_config_filter(config, config_filter) return config diff --git a/monkey/monkey_island/cc/services/mode/set_island_mode_service.py b/monkey/monkey_island/cc/services/mode/set_island_mode_service.py index b11aab8ca..dd70070d1 100644 --- a/monkey/monkey_island/cc/services/mode/set_island_mode_service.py +++ b/monkey/monkey_island/cc/services/mode/set_island_mode_service.py @@ -7,4 +7,4 @@ def set_mode(mode: IslandModeEnum): island_mode_model = IslandMode() island_mode_model.mode = mode.value island_mode_model.save() - ConfigService.update_config_on_mode_set(mode.value) + ConfigService.update_config_on_mode_set(mode)