Move database reset to happen during the registration

This commit is contained in:
VakarisZ 2021-09-29 16:41:08 +03:00
parent b73958dd55
commit c211d51d8c
4 changed files with 12 additions and 28 deletions

View File

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

View File

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

View File

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

View File

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