forked from p15670423/monkey
Backmerged and fixed a bug where mitigations got deleted on DB reset
This commit is contained in:
parent
bc166ef67d
commit
7adbbdd525
|
@ -9,6 +9,8 @@ from monkey_island.cc.services.attack.test_mitre_api_interface import MitreApiIn
|
||||||
|
|
||||||
class AttackMitigations(Document):
|
class AttackMitigations(Document):
|
||||||
|
|
||||||
|
COLLECTION_NAME = "attack_mitigations"
|
||||||
|
|
||||||
technique_id = StringField(required=True, primary_key=True)
|
technique_id = StringField(required=True, primary_key=True)
|
||||||
mitigations = ListField(EmbeddedDocumentField('Mitigation'))
|
mitigations = ListField(EmbeddedDocumentField('Mitigation'))
|
||||||
|
|
||||||
|
|
|
@ -3,6 +3,7 @@ import logging
|
||||||
from monkey_island.cc.services.config import ConfigService
|
from monkey_island.cc.services.config import ConfigService
|
||||||
from monkey_island.cc.services.attack.attack_config import AttackConfig
|
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.services.post_breach_files import remove_PBA_files
|
||||||
|
from monkey_island.cc.models.attack.attack_mitigations import AttackMitigations
|
||||||
from flask import jsonify
|
from flask import jsonify
|
||||||
from monkey_island.cc.database import mongo
|
from monkey_island.cc.database import mongo
|
||||||
|
|
||||||
|
@ -17,7 +18,8 @@ class Database(object):
|
||||||
def reset_db():
|
def reset_db():
|
||||||
remove_PBA_files()
|
remove_PBA_files()
|
||||||
# We can't drop system collections.
|
# 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()
|
ConfigService.init_config()
|
||||||
AttackConfig.reset_config()
|
AttackConfig.reset_config()
|
||||||
logger.info('DB was reset')
|
logger.info('DB was reset')
|
||||||
|
|
|
@ -9,7 +9,7 @@ def setup():
|
||||||
|
|
||||||
|
|
||||||
def try_store_mitigations_on_mongo():
|
def try_store_mitigations_on_mongo():
|
||||||
mitigation_collection_name = 'attack_mitigations'
|
mitigation_collection_name = AttackMitigations.COLLECTION_NAME
|
||||||
try:
|
try:
|
||||||
mongo.db.validate_collection(mitigation_collection_name)
|
mongo.db.validate_collection(mitigation_collection_name)
|
||||||
if mongo.db.attack_mitigations.count() == 0:
|
if mongo.db.attack_mitigations.count() == 0:
|
||||||
|
|
Loading…
Reference in New Issue