Improved performance by storing mitigations on mongodb

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:

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

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

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

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

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

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

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

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

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
 

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 {}
 

monkey/monkey_island/cc/services/attack/technique_reports/attack_mitigations.py

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

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)