From 57df09986308cae17af222676e2a17cb60b9b801 Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Thu, 26 Mar 2020 12:58:41 +0200 Subject: [PATCH] Improved performance by storing mitigations on mongodb --- monkey/monkey_island/cc/main.py | 3 +++ .../cc/models/attack_mitigation.py | 23 ++++++++++++++++++ .../cc/services/attack/attack_report.py | 1 - .../cc/services/attack/mitre_api_interface.py | 13 ++++++++++ .../attack/technique_reports/T1003.py | 2 +- .../attack/technique_reports/T1059.py | 2 +- .../attack/technique_reports/T1075.py | 2 +- .../attack/technique_reports/T1082.py | 2 +- .../attack/technique_reports/T1086.py | 2 +- .../attack/technique_reports/T1210.py | 2 +- .../attack/technique_reports/__init__.py | 12 ++++++---- .../technique_reports/attack_mitigations.py | 15 ------------ monkey/monkey_island/cc/setup.py | 24 +++++++++++++++++++ 13 files changed, 76 insertions(+), 27 deletions(-) create mode 100644 monkey/monkey_island/cc/models/attack_mitigation.py create mode 100644 monkey/monkey_island/cc/services/attack/mitre_api_interface.py delete mode 100644 monkey/monkey_island/cc/services/attack/technique_reports/attack_mitigations.py create mode 100644 monkey/monkey_island/cc/setup.py diff --git a/monkey/monkey_island/cc/main.py b/monkey/monkey_island/cc/main.py index f06d36ea3..045f5452e 100644 --- a/monkey/monkey_island/cc/main.py +++ b/monkey/monkey_island/cc/main.py @@ -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: diff --git a/monkey/monkey_island/cc/models/attack_mitigation.py b/monkey/monkey_island/cc/models/attack_mitigation.py new file mode 100644 index 000000000..632795bb4 --- /dev/null +++ b/monkey/monkey_island/cc/models/attack_mitigation.py @@ -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() diff --git a/monkey/monkey_island/cc/services/attack/attack_report.py b/monkey/monkey_island/cc/services/attack/attack_report.py index 7184ce202..db80c76c6 100644 --- a/monkey/monkey_island/cc/services/attack/attack_report.py +++ b/monkey/monkey_island/cc/services/attack/attack_report.py @@ -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() diff --git a/monkey/monkey_island/cc/services/attack/mitre_api_interface.py b/monkey/monkey_island/cc/services/attack/mitre_api_interface.py new file mode 100644 index 000000000..b1da945ca --- /dev/null +++ b/monkey/monkey_island/cc/services/attack/mitre_api_interface.py @@ -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 diff --git a/monkey/monkey_island/cc/services/attack/technique_reports/T1003.py b/monkey/monkey_island/cc/services/attack/technique_reports/T1003.py index cba869432..cbd3bf8bf 100644 --- a/monkey/monkey_island/cc/services/attack/technique_reports/T1003.py +++ b/monkey/monkey_island/cc/services/attack/technique_reports/T1003.py @@ -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 diff --git a/monkey/monkey_island/cc/services/attack/technique_reports/T1059.py b/monkey/monkey_island/cc/services/attack/technique_reports/T1059.py index 2cef70134..c0e4dc3f6 100644 --- a/monkey/monkey_island/cc/services/attack/technique_reports/T1059.py +++ b/monkey/monkey_island/cc/services/attack/technique_reports/T1059.py @@ -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 diff --git a/monkey/monkey_island/cc/services/attack/technique_reports/T1075.py b/monkey/monkey_island/cc/services/attack/technique_reports/T1075.py index a9776b136..370db3ca2 100644 --- a/monkey/monkey_island/cc/services/attack/technique_reports/T1075.py +++ b/monkey/monkey_island/cc/services/attack/technique_reports/T1075.py @@ -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 diff --git a/monkey/monkey_island/cc/services/attack/technique_reports/T1082.py b/monkey/monkey_island/cc/services/attack/technique_reports/T1082.py index ce91f0b16..cdbdb42ec 100644 --- a/monkey/monkey_island/cc/services/attack/technique_reports/T1082.py +++ b/monkey/monkey_island/cc/services/attack/technique_reports/T1082.py @@ -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 diff --git a/monkey/monkey_island/cc/services/attack/technique_reports/T1086.py b/monkey/monkey_island/cc/services/attack/technique_reports/T1086.py index b46afb479..897ccdaaf 100644 --- a/monkey/monkey_island/cc/services/attack/technique_reports/T1086.py +++ b/monkey/monkey_island/cc/services/attack/technique_reports/T1086.py @@ -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 diff --git a/monkey/monkey_island/cc/services/attack/technique_reports/T1210.py b/monkey/monkey_island/cc/services/attack/technique_reports/T1210.py index b6e606be2..babe5c788 100644 --- a/monkey/monkey_island/cc/services/attack/technique_reports/T1210.py +++ b/monkey/monkey_island/cc/services/attack/technique_reports/T1210.py @@ -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 diff --git a/monkey/monkey_island/cc/services/attack/technique_reports/__init__.py b/monkey/monkey_island/cc/services/attack/technique_reports/__init__.py index dcb3552b3..f620996fb 100644 --- a/monkey/monkey_island/cc/services/attack/technique_reports/__init__.py +++ b/monkey/monkey_island/cc/services/attack/technique_reports/__init__.py @@ -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 {} diff --git a/monkey/monkey_island/cc/services/attack/technique_reports/attack_mitigations.py b/monkey/monkey_island/cc/services/attack/technique_reports/attack_mitigations.py deleted file mode 100644 index ffd79a4d3..000000000 --- a/monkey/monkey_island/cc/services/attack/technique_reports/attack_mitigations.py +++ /dev/null @@ -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 diff --git a/monkey/monkey_island/cc/setup.py b/monkey/monkey_island/cc/setup.py new file mode 100644 index 000000000..9a501d609 --- /dev/null +++ b/monkey/monkey_island/cc/setup.py @@ -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)