Move database reset to happen during the registration

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)

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")

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())

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():