From c211d51d8caa29de4a7a5c8d9ed98d9cafd71b2c Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Wed, 29 Sep 2021 16:41:08 +0300 Subject: [PATCH] Move database reset to happen during the registration --- monkey/monkey_island/cc/resources/auth/auth.py | 9 --------- .../cc/resources/auth/registration.py | 18 ++---------------- monkey/monkey_island/cc/services/database.py | 4 ++++ .../cc/setup/mongo/database_initializer.py | 9 ++++++--- 4 files changed, 12 insertions(+), 28 deletions(-) diff --git a/monkey/monkey_island/cc/resources/auth/auth.py b/monkey/monkey_island/cc/resources/auth/auth.py index 98408c05c..064395eaf 100644 --- a/monkey/monkey_island/cc/resources/auth/auth.py +++ b/monkey/monkey_island/cc/resources/auth/auth.py @@ -11,9 +11,6 @@ from jwt import PyJWTError import monkey_island.cc.environment.environment_singleton as env_singleton import monkey_island.cc.resources.auth.password_utils as password_utils import monkey_island.cc.resources.auth.user_store as user_store -from monkey_island.cc.database import mongo -from monkey_island.cc.models.attack.attack_mitigations import AttackMitigations -from monkey_island.cc.setup.mongo.database_initializer import init_collections logger = logging.getLogger(__name__) @@ -45,7 +42,6 @@ class Authenticate(flask_restful.Resource): if _credentials_match_registered_user(username, password): access_token = _create_access_token(username) - _check_attack_mitigations_in_mongo() return make_response({"access_token": access_token, "error": ""}, 200) else: return make_response({"error": "Invalid credentials"}, 401) @@ -78,11 +74,6 @@ def _create_access_token(username): return access_token -def _check_attack_mitigations_in_mongo(): - if AttackMitigations.COLLECTION_NAME not in mongo.db.list_collection_names(): - init_collections() - - # See https://flask-jwt-extended.readthedocs.io/en/stable/custom_decorators/ def jwt_required(fn): @wraps(fn) diff --git a/monkey/monkey_island/cc/resources/auth/registration.py b/monkey/monkey_island/cc/resources/auth/registration.py index 92fca24c9..12c17d6e5 100644 --- a/monkey/monkey_island/cc/resources/auth/registration.py +++ b/monkey/monkey_island/cc/resources/auth/registration.py @@ -7,9 +7,8 @@ from flask import make_response, request import monkey_island.cc.environment.environment_singleton as env_singleton import monkey_island.cc.resources.auth.password_utils as password_utils from common.utils.exceptions import InvalidRegistrationCredentialsError, RegistrationNotNeededError -from monkey_island.cc.database import mongo from monkey_island.cc.environment.user_creds import UserCreds -from monkey_island.cc.setup.mongo.database_initializer import init_collections +from monkey_island.cc.setup.mongo.database_initializer import reset_database logger = logging.getLogger(__name__) @@ -17,9 +16,6 @@ logger = logging.getLogger(__name__) class Registration(flask_restful.Resource): def get(self): is_registration_needed = env_singleton.env.needs_registration() - if is_registration_needed: - # if registration is required, drop previous user's data (for credentials reset case) - _drop_mongo_db() return {"needs_registration": is_registration_needed} def post(self): @@ -27,16 +23,10 @@ class Registration(flask_restful.Resource): try: env_singleton.env.try_add_user(credentials) - init_collections() + reset_database() return make_response({"error": ""}, 200) except (InvalidRegistrationCredentialsError, RegistrationNotNeededError) as e: return make_response({"error": str(e)}, 400) - except Exception as ex: - logger.error( - "Exception raised during registration; most likely an issue with the " - f"mongo collection's initialisation. Exception: {str(ex)}." - ) - return make_response({"error": str(ex)}, 400) def _get_user_credentials_from_request(request): @@ -47,7 +37,3 @@ def _get_user_credentials_from_request(request): password_hash = password_utils.hash_password(password) return UserCreds(username, password_hash) - - -def _drop_mongo_db(): - mongo.db.command("dropDatabase") diff --git a/monkey/monkey_island/cc/services/database.py b/monkey/monkey_island/cc/services/database.py index d0656f946..afd4ecc02 100644 --- a/monkey/monkey_island/cc/services/database.py +++ b/monkey/monkey_island/cc/services/database.py @@ -37,3 +37,7 @@ class Database(object): def init_db(): if not mongo.db.collection_names(): Database.reset_db() + + @staticmethod + def is_mitigations_missing() -> bool: + return bool(AttackMitigations.COLLECTION_NAME not in mongo.db.list_collection_names()) diff --git a/monkey/monkey_island/cc/setup/mongo/database_initializer.py b/monkey/monkey_island/cc/setup/mongo/database_initializer.py index 34914c7ce..72d7bec7d 100644 --- a/monkey/monkey_island/cc/setup/mongo/database_initializer.py +++ b/monkey/monkey_island/cc/setup/mongo/database_initializer.py @@ -5,13 +5,16 @@ from pymongo import errors from monkey_island.cc.database import mongo from monkey_island.cc.models.attack.attack_mitigations import AttackMitigations from monkey_island.cc.services.attack.mitre_api_interface import MitreApiInterface +from monkey_island.cc.services.database import Database logger = logging.getLogger(__name__) -def init_collections(): - logger.info("Setting up the Monkey Island, this might take a while...") - _try_store_mitigations_on_mongo() +def reset_database(): + Database.reset_db() + if Database.is_mitigations_missing(): + logger.info("Populating Monkey Island with ATT&CK mitigations, this might take a while...") + _try_store_mitigations_on_mongo() def _try_store_mitigations_on_mongo():