Improved performance by storing mitigations on mongodb

This commit is contained in:
VakarisZ 2020-03-26 12:58:41 +02:00
parent 06e8156a4a
commit 57df099863
13 changed files with 76 additions and 27 deletions

View File

@ -26,6 +26,7 @@ from monkey_island.cc.environment.environment import env
from monkey_island.cc.database import is_db_server_up, get_db_version
from monkey_island.cc.resources.monkey_download import MonkeyDownload
from common.version import get_version
from monkey_island.cc.setup import setup
def main():
@ -43,6 +44,8 @@ def main():
crt_path = os.path.join(MONKEY_ISLAND_ABS_PATH, 'cc', 'server.crt')
key_path = os.path.join(MONKEY_ISLAND_ABS_PATH, 'cc', 'server.key')
setup()
if env.is_debug():
app.run(host='0.0.0.0', debug=True, ssl_context=(crt_path, key_path))
else:

View File

@ -0,0 +1,23 @@
from mongoengine import Document, StringField, DoesNotExist
class AttackMitigation(Document):
technique_id = StringField(required=True, primary_key=True)
name = StringField(required=True)
description = StringField(required=True)
@staticmethod
def get_mitigation_by_technique_id(technique_id: str) -> Document:
try:
return AttackMitigation.objects.get(technique_id=technique_id)
except DoesNotExist:
raise Exception("Attack technique with id {} does not exist!".format(technique_id))
@staticmethod
def add_mitigation_from_stix2(mitigation_stix2_data):
mitigation_model = AttackMitigation(technique_id=mitigation_stix2_data['external_references'][0]['external_id'],
name=mitigation_stix2_data['name'],
description=mitigation_stix2_data['description'])
if mitigation_model.technique_id.startswith('T'):
mitigation_model.save()

View File

@ -57,7 +57,6 @@ class AttackReportService:
'meta': {'latest_monkey_modifytime': Monkey.get_latest_modifytime()},
'name': REPORT_NAME
}
for tech_id, tech_info in list(AttackConfig.get_techniques_for_report().items()):
try:
technique_report_data = TECHNIQUES[tech_id].get_report_data()

View File

@ -0,0 +1,13 @@
from stix2 import FileSystemSource, Filter
class MitreApiInterface:
ATTACK_DATA_PATH = 'monkey_island/cc/services/attack/attack_data/enterprise-attack'
@staticmethod
def get_all_mitigations() -> list:
file_system = FileSystemSource(MitreApiInterface.ATTACK_DATA_PATH)
mitigation_filter = [Filter('type', '=', 'course-of-action')]
all_mitigations = file_system.query(mitigation_filter)
return all_mitigations

View File

@ -24,7 +24,7 @@ class T1003(AttackTechnique):
else:
status = ScanStatus.UNSCANNED.value
data.update(T1003.get_message_and_status(status))
data.update(T1003.get_mitigations_by_status(status))
data.update(T1003.get_mitigation_by_status(status))
data['stolen_creds'] = ReportService.get_stolen_creds()
data['stolen_creds'].extend(ReportService.get_ssh_keys())
return data

View File

@ -30,5 +30,5 @@ class T1059(AttackTechnique):
else:
status = ScanStatus.UNSCANNED.value
data.update(T1059.get_message_and_status(status))
data.update(T1059.get_mitigations_by_status(status))
data.update(T1059.get_mitigation_by_status(status))
return data

View File

@ -40,5 +40,5 @@ class T1075(AttackTechnique):
else:
status = ScanStatus.UNSCANNED.value
data.update(T1075.get_message_and_status(status))
data.update(T1075.get_mitigations_by_status(status))
data.update(T1075.get_mitigation_by_status(status))
return data

View File

@ -44,6 +44,6 @@ class T1082(AttackTechnique):
status = ScanStatus.USED.value
else:
status = ScanStatus.UNSCANNED.value
data.update(T1082.get_mitigations_by_status(status))
data.update(T1082.get_mitigation_by_status(status))
data.update(T1082.get_message_and_status(status))
return data

View File

@ -32,6 +32,6 @@ class T1086(AttackTechnique):
else:
status = ScanStatus.UNSCANNED.value
data.update(T1086.get_mitigations_by_status(status))
data.update(T1086.get_mitigation_by_status(status))
data.update(T1086.get_message_and_status(status))
return data

View File

@ -23,7 +23,7 @@ class T1210(AttackTechnique):
else:
status = ScanStatus.UNSCANNED.value
data.update(T1210.get_message_and_status(status))
data.update(T1210.get_mitigations_by_status(status))
data.update(T1210.get_mitigation_by_status(status))
data.update({'scanned_services': scanned_services, 'exploited_services': exploited_services})
return data

View File

@ -5,7 +5,7 @@ from monkey_island.cc.database import mongo
from common.utils.attack_utils import ScanStatus
from monkey_island.cc.services.attack.attack_config import AttackConfig
from common.utils.code_utils import abstractstatic
from monkey_island.cc.services.attack.technique_reports.attack_mitigations import AttackMitigations
from monkey_island.cc.models.attack_mitigation import AttackMitigation
logger = logging.getLogger(__name__)
@ -112,20 +112,22 @@ class AttackTechnique(object, metaclass=abc.ABCMeta):
data.update({'status': status,
'title': title,
'message': cls.get_message_by_status(status)})
data.update(cls.get_mitigations_by_status(status))
data.update(cls.get_mitigation_by_status(status))
return data
@classmethod
def get_base_data_by_status(cls, status):
data = cls.get_message_and_status(status)
data.update({'title': cls.technique_title()})
data.update(cls.get_mitigations_by_status(status))
data.update(cls.get_mitigation_by_status(status))
return data
@classmethod
def get_mitigations_by_status(cls, status: ScanStatus) -> dict:
def get_mitigation_by_status(cls, status: ScanStatus) -> dict:
if status == ScanStatus.USED.value:
return AttackMitigations.get_mitigations_by_id(cls.tech_id)
mitigation_document = AttackMitigation.get_mitigation_by_technique_id(str(cls.tech_id))
return {'mitigations': {'name': mitigation_document['name'],
'description': mitigation_document['description']}}
else:
return {}

View File

@ -1,15 +0,0 @@
from stix2 import FileSystemSource, Filter, parse
class AttackMitigations:
@staticmethod
def get_mitigations_by_id(technique_id: str) -> dict:
file_system = FileSystemSource('monkey_island/cc/services/attack/attack_data/enterprise-attack')
technique_filter = [
Filter('type', '=', 'course-of-action'),
Filter('external_references.external_id', '=', str(technique_id))
]
mitigations = parse(file_system.query(technique_filter)[0], allow_custom=True)
mitigations = {'mitigations': {'description': mitigations['description'], 'name': mitigations['name']}}
return mitigations

View File

@ -0,0 +1,24 @@
from monkey_island.cc.services.attack.mitre_api_interface import MitreApiInterface
from monkey_island.cc.models.attack_mitigation import AttackMitigation
from monkey_island.cc.database import mongo
from pymongo import errors
def setup():
try_store_mitigations_on_mongo()
def try_store_mitigations_on_mongo():
# import the 'errors' module from PyMongo
mitigation_collection_name = 'attack_mitigation'
try:
mongo.db.validate_collection(mitigation_collection_name)
except errors.OperationFailure:
mongo.db.create_collection(mitigation_collection_name)
store_mitigations_on_mongo()
def store_mitigations_on_mongo():
all_mitigations = MitreApiInterface.get_all_mitigations()
for mitigation in all_mitigations:
AttackMitigation.add_mitigation_from_stix2(mitigation)