island: On login, check if collection 'attack_mitigations' is present in DB,

add if not
This commit is contained in:
Shreya Malviya 2021-09-27 13:52:21 +05:30 committed by VakarisZ
parent 3cbeb3dbf7
commit 194e244080
1 changed files with 9 additions and 0 deletions

View File

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