island: Pass `mode` as an enum instead of enum value when updating config, and update typehints

This commit is contained in:
Shreya 2021-07-20 13:16:06 +05:30
parent 5603773083
commit 0681782505
2 changed files with 6 additions and 5 deletions

View File

@ -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_filters import FILTER_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.get_island_mode_service import ModeNotSetError, get_mode 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.post_breach_files import PostBreachFilesService
from monkey_island.cc.services.utils.network_utils import local_ip_addresses from monkey_island.cc.services.utils.network_utils import local_ip_addresses
@ -248,18 +249,18 @@ class ConfigService:
return config return config
@staticmethod @staticmethod
def update_config_on_mode_set(mode: str) -> None: def update_config_on_mode_set(mode: IslandModeEnum) -> None:
config = ConfigService.get_config() config = ConfigService.get_config()
ConfigService.update_config_per_mode(mode, config, True) ConfigService.update_config_per_mode(mode, config, True)
@staticmethod @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) config = ConfigService._set_default_config_values_per_mode(mode, config)
ConfigService.update_config(config_json=config, should_encrypt=True) ConfigService.update_config(config_json=config, should_encrypt=True)
@staticmethod @staticmethod
def _set_default_config_values_per_mode(mode: str, config: Dict) -> Dict: def _set_default_config_values_per_mode(mode: IslandModeEnum, config: Dict) -> Dict:
config_filter = FILTER_PER_MODE[mode] config_filter = FILTER_PER_MODE[mode.value]
config = ConfigService._apply_config_filter(config, config_filter) config = ConfigService._apply_config_filter(config, config_filter)
return config return config

View File

@ -7,4 +7,4 @@ def set_mode(mode: IslandModeEnum):
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()
ConfigService.update_config_on_mode_set(mode.value) ConfigService.update_config_on_mode_set(mode)