From 7adbbdd525c5576f7fd133db610ec59b88fd777c Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Wed, 1 Apr 2020 09:21:48 +0300 Subject: [PATCH] Backmerged and fixed a bug where mitigations got deleted on DB reset --- monkey/monkey_island/cc/models/attack/attack_mitigations.py | 2 ++ monkey/monkey_island/cc/services/database.py | 4 +++- monkey/monkey_island/cc/setup.py | 2 +- 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/models/attack/attack_mitigations.py b/monkey/monkey_island/cc/models/attack/attack_mitigations.py index 8b2851268..a51778c48 100644 --- a/monkey/monkey_island/cc/models/attack/attack_mitigations.py +++ b/monkey/monkey_island/cc/models/attack/attack_mitigations.py @@ -9,6 +9,8 @@ from monkey_island.cc.services.attack.test_mitre_api_interface import MitreApiIn class AttackMitigations(Document): + COLLECTION_NAME = "attack_mitigations" + technique_id = StringField(required=True, primary_key=True) mitigations = ListField(EmbeddedDocumentField('Mitigation')) diff --git a/monkey/monkey_island/cc/services/database.py b/monkey/monkey_island/cc/services/database.py index 7062d71c3..5961fa70b 100644 --- a/monkey/monkey_island/cc/services/database.py +++ b/monkey/monkey_island/cc/services/database.py @@ -3,6 +3,7 @@ import logging from monkey_island.cc.services.config import ConfigService from monkey_island.cc.services.attack.attack_config import AttackConfig from monkey_island.cc.services.post_breach_files import remove_PBA_files +from monkey_island.cc.models.attack.attack_mitigations import AttackMitigations from flask import jsonify from monkey_island.cc.database import mongo @@ -17,7 +18,8 @@ class Database(object): def reset_db(): remove_PBA_files() # We can't drop system collections. - [mongo.db[x].drop() for x in mongo.db.collection_names() if not x.startswith('system.')] + [mongo.db[x].drop() for x in mongo.db.collection_names() if not x.startswith('system.') + and not AttackMitigations.COLLECTION_NAME] ConfigService.init_config() AttackConfig.reset_config() logger.info('DB was reset') diff --git a/monkey/monkey_island/cc/setup.py b/monkey/monkey_island/cc/setup.py index e1b53affa..c37c24f3b 100644 --- a/monkey/monkey_island/cc/setup.py +++ b/monkey/monkey_island/cc/setup.py @@ -9,7 +9,7 @@ def setup(): def try_store_mitigations_on_mongo(): - mitigation_collection_name = 'attack_mitigations' + mitigation_collection_name = AttackMitigations.COLLECTION_NAME try: mongo.db.validate_collection(mitigation_collection_name) if mongo.db.attack_mitigations.count() == 0: