From 194e244080ba2ac74afcfd78b6912bc63c41b6a6 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 27 Sep 2021 13:52:21 +0530 Subject: [PATCH] island: On login, check if collection 'attack_mitigations' is present in DB, add if not --- monkey/monkey_island/cc/resources/auth/auth.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/monkey/monkey_island/cc/resources/auth/auth.py b/monkey/monkey_island/cc/resources/auth/auth.py index 064395eaf..98408c05c 100644 --- a/monkey/monkey_island/cc/resources/auth/auth.py +++ b/monkey/monkey_island/cc/resources/auth/auth.py @@ -11,6 +11,9 @@ 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__) @@ -42,6 +45,7 @@ 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) @@ -74,6 +78,11 @@ 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)