diff --git a/monkey/monkey_island/cc/resources/configuration_import.py b/monkey/monkey_island/cc/resources/configuration_import.py index 549ff1a93..40413c3f5 100644 --- a/monkey/monkey_island/cc/resources/configuration_import.py +++ b/monkey/monkey_island/cc/resources/configuration_import.py @@ -15,20 +15,27 @@ from monkey_island.cc.resources.auth.auth import jwt_required from monkey_island.cc.services.config import ConfigService from monkey_island.cc.services.utils.config_encryption import decrypt_config +logger = logging.getLogger(__name__) + + +class ImportStatuses: + UNSAFE_OPTION_VERIFICATION_REQUIRED = "unsafe_options_verification_required" + INVALID_CONFIGURATION = "invalid_configuration" + PASSWORD_REQUIRED = "password_required" + WRONG_PASSWORD = "wrong_password" + IMPORTED = "imported" + @dataclass class ResponseContents: - import_status: str = "imported" + import_status: str = ImportStatuses.IMPORTED message: str = "" status_code: int = 200 config: str = "" config_schema: str = "" def form_response(self): - return self.__dict__, self.status_code - - -logger = logging.getLogger(__name__) + return self.__dict__ class ConfigurationImport(flask_restful.Resource): @@ -38,11 +45,7 @@ class ConfigurationImport(flask_restful.Resource): def post(self): request_contents = json.loads(request.data) try: - try: - config = json.loads(request_contents["config"]) - except JSONDecodeError: - config = decrypt_config(request_contents["config"], request_contents["password"]) - + config = ConfigurationImport._get_plaintext_config_from_request(request_contents) if request_contents["unsafeOptionsVerified"]: ConfigurationImport.import_config(config) return ResponseContents().form_response() @@ -50,28 +53,31 @@ class ConfigurationImport(flask_restful.Resource): return ResponseContents( config=json.dumps(config), config_schema=ConfigService.get_config_schema(), - import_status="unsafe_options_verification_required", - status_code=403, + import_status=ImportStatuses.UNSAFE_OPTION_VERIFICATION_REQUIRED, ).form_response() except InvalidCredentialsError: return ResponseContents( - import_status="wrong_password", message="Wrong password supplied", status_code=403 + import_status=ImportStatuses.WRONG_PASSWORD, message="Wrong password supplied" ).form_response() except InvalidConfigurationError: return ResponseContents( - import_status="invalid_configuration", + import_status=ImportStatuses.INVALID_CONFIGURATION, message="Invalid configuration supplied. " "Maybe the format is outdated or the file is corrupted.", - status_code=400, ).form_response() except NoCredentialsError: return ResponseContents( - import_status="password_required", - message="Configuration file is protected with a password. " - "Please enter the password", - status_code=403, + import_status=ImportStatuses.PASSWORD_REQUIRED, ).form_response() + @staticmethod + def _get_plaintext_config_from_request(request_contents: dict) -> dict: + try: + config = json.loads(request_contents["config"]) + except JSONDecodeError: + config = decrypt_config(request_contents["config"], request_contents["password"]) + return config + @staticmethod def import_config(config_json): if not ConfigService.update_config(config_json, should_encrypt=True):