Improved readability in configuration_import.py by removing unused variables and extracting methods.

This commit is contained in:
VakarisZ 2021-06-03 16:56:43 +03:00
parent c25ea0edf8
commit 2f9c6bf035
1 changed files with 25 additions and 19 deletions

View File

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