Refactoring ZT findings

monkey/monkey_island/cc/models/zero_trust/aggregate_finding.py

@@ -1,5 +1,9 @@
+from typing import List
+
 import common.common_consts.zero_trust_consts as zero_trust_consts
+from monkey_island.cc.models.zero_trust.event import Event
 from monkey_island.cc.models.zero_trust.finding import Finding
+from monkey_island.cc.models.zero_trust.finding_details import FindingDetails
 
 
 class AggregateFinding(Finding):
@@ -12,15 +16,25 @@ class AggregateFinding(Finding):
         :raises: Assertion error if this is used when there's more then one finding which fits the query - this is not
         when this function should be used.
         """
-        existing_findings = Finding.objects(test=test, status=status).exclude('events')
+        existing_findings = Finding.objects(test=test, status=status)
         assert (len(existing_findings) < 2), "More than one finding exists for {}:{}".format(test, status)
 
         if len(existing_findings) == 0:
-            Finding.save_finding(test, status, events)
+            AggregateFinding.create_new_finding(test, status, events)
         else:
             # Now we know for sure this is the only one
-            orig_finding = existing_findings[0]
+            AggregateFinding.add_events(existing_findings[0], events)
-            orig_finding.add_events(events)
+
+    @staticmethod
+    def create_new_finding(test: str, status: str, events: List[Event]):
+        details = FindingDetails()
+        details.events = events
+        details.save()
+        Finding.save_finding(test, status, details)
+
+    @staticmethod
+    def add_events(finding: Finding, events: List[Event]):
+        finding.details.fetch().add_events(events)
 
 
 def add_malicious_activity_to_timeline(events):

monkey/monkey_island/cc/models/zero_trust/finding.py

@@ -4,12 +4,13 @@ Define a Document Schema for Zero Trust findings.
 """
 from typing import List
 
-from mongoengine import Document, EmbeddedDocumentListField, StringField
+from mongoengine import Document, EmbeddedDocumentListField, StringField, LazyReferenceField
 
 import common.common_consts.zero_trust_consts as zero_trust_consts
 # Dummy import for mongoengine.
 # noinspection PyUnresolvedReferences
 from monkey_island.cc.models.zero_trust.event import Event
+from monkey_island.cc.models.zero_trust.finding_details import FindingDetails
 
 
 class Finding(Document):
@@ -33,7 +34,7 @@ class Finding(Document):
     # SCHEMA
     test = StringField(required=True, choices=zero_trust_consts.TESTS)
     status = StringField(required=True, choices=zero_trust_consts.ORDERED_TEST_STATUSES)
-    events = EmbeddedDocumentListField(document_type=Event)
+    details = LazyReferenceField(document_type=FindingDetails, required=True)
     # http://docs.mongoengine.org/guide/defining-documents.html#document-inheritance
     meta = {'allow_inheritance': True}
 
@@ -46,15 +47,11 @@ class Finding(Document):
 
     # Creation methods
     @staticmethod
-    def save_finding(test, status, events):
+    def save_finding(test: str, status: str, detail_ref):
-        finding = Finding(
+        finding = Finding(test=test,
-            test=test,
+                          status=status,
-            status=status,
+                          details=detail_ref)
-            events=events)
 
         finding.save()
 
         return finding
-
-    def add_events(self, events: List) -> None:
-        self.update(push_all__events=events)

monkey/monkey_island/cc/models/zero_trust/finding_details.py

@@ -0,0 +1,27 @@
+from datetime import datetime
+from typing import List
+
+from mongoengine import DateTimeField, Document, StringField, EmbeddedDocumentListField
+
+import common.common_consts.zero_trust_consts as zero_trust_consts
+from monkey_island.cc.models.zero_trust.event import Event
+from monkey_island.cc.models.zero_trust.scoutsuite_finding import ScoutsuiteFinding
+
+
+class FindingDetails(Document):
+    """
+    This model represents additional information about monkey finding:
+    Events if monkey finding
+    Scoutsuite findings if scoutsuite finding
+    """
+
+    # SCHEMA
+    events = EmbeddedDocumentListField(document_type=Event, required=False)
+    scoutsuite_findings = EmbeddedDocumentListField(document_type=ScoutsuiteFinding, required=False)
+
+    # LOGIC
+    def add_events(self, events: List[Event]) -> None:
+        self.update(push_all__events=events)
+
+    def add_scoutsuite_findings(self, scoutsuite_findings: List[ScoutsuiteFinding]) -> None:
+        self.update(push_all__scoutsuite_findings=scoutsuite_findings)

monkey/monkey_island/cc/models/zero_trust/scoutsuite_finding.py

@@ -0,0 +1,17 @@
+from datetime import datetime
+
+from mongoengine import DateTimeField, EmbeddedDocument, StringField
+
+
+class ScoutsuiteFinding(EmbeddedDocument):
+    # SCHEMA
+    temp = StringField(required=True)
+
+    # LOGIC
+    @staticmethod
+    def create_scoutsuite_finding(title, message, event_type, timestamp=None):
+        scoutsuite_finding = ScoutsuiteFinding()
+
+        scoutsuite_finding.temp = "temp"
+
+        return scoutsuite_finding

monkey/monkey_island/cc/resources/reporting/report.py

@@ -35,6 +35,6 @@ class Report(flask_restful.Resource):
             elif report_data == REPORT_DATA_PRINCIPLES_STATUS:
                 return jsonify(ZeroTrustService.get_principles_status())
             elif report_data == REPORT_DATA_FINDINGS:
-                return jsonify(ZeroTrustService.get_all_findings())
+                return jsonify(ZeroTrustService.get_all_monkey_findings())
 
         flask_restful.abort(http.client.NOT_FOUND)

monkey/monkey_island/cc/resources/telemetry.py

@@ -46,6 +46,7 @@ class Telemetry(flask_restful.Resource):
     @TestTelemStore.store_test_telem
     def post(self):
         telemetry_json = json.loads(request.data)
+        telemetry_json['data'] = json.loads(telemetry_json['data'])
         telemetry_json['timestamp'] = datetime.now()
         telemetry_json['command_control_channel'] = {'src': request.remote_addr, 'dst': request.host}
 

monkey/monkey_island/cc/services/reporting/zero_trust_service.py

@@ -7,6 +7,8 @@ from monkey_island.cc.models.zero_trust.finding import Finding
 
 # How many events of a single finding to return to UI.
 # 50 will return 50 latest and 50 oldest events from a finding
+from monkey_island.cc.models.zero_trust.finding_details import FindingDetails
+
 EVENT_FETCH_CNT = 50
 
 
@@ -14,7 +16,7 @@ class ZeroTrustService(object):
     @staticmethod
     def get_pillars_grades():
         pillars_grades = []
-        all_findings = Finding.objects().exclude('events')
+        all_findings = Finding.objects()
         for pillar in zero_trust_consts.PILLARS:
             pillars_grades.append(ZeroTrustService.__get_pillar_grade(pillar, all_findings))
         return pillars_grades
@@ -41,7 +43,8 @@ class ZeroTrustService(object):
             if pillar in test_info[zero_trust_consts.PILLARS_KEY]:
                 pillar_grade[finding.status] += 1
 
-        pillar_grade[zero_trust_consts.STATUS_UNEXECUTED] = sum(1 for condition in list(test_unexecuted.values()) if condition)
+        pillar_grade[zero_trust_consts.STATUS_UNEXECUTED] = sum(1 for condition in
+                                                                list(test_unexecuted.values()) if condition)
 
         return pillar_grade
 
@@ -70,7 +73,7 @@ class ZeroTrustService(object):
         worst_status = zero_trust_consts.STATUS_UNEXECUTED
         all_statuses = set()
         for test in principle_tests:
-            all_statuses |= set(Finding.objects(test=test).exclude('events').distinct("status"))
+            all_statuses |= set(Finding.objects(test=test).distinct('status'))
 
         for status in all_statuses:
             if zero_trust_consts.ORDERED_TEST_STATUSES.index(status) \
@@ -83,7 +86,7 @@ class ZeroTrustService(object):
     def __get_tests_status(principle_tests):
         results = []
         for test in principle_tests:
-            test_findings = Finding.objects(test=test).exclude('events')
+            test_findings = Finding.objects(test=test)
             results.append(
                 {
                     "test": zero_trust_consts.TESTS_MAP[test][zero_trust_consts.TEST_EXPLANATION_KEY],
@@ -108,18 +111,30 @@ class ZeroTrustService(object):
         return current_worst_status
 
     @staticmethod
-    def get_all_findings():
+    def get_all_monkey_findings():
+        findings = list(Finding.objects)
+        for finding in findings:
+            details = finding.details.fetch()
+            finding.details = details
+        enriched_findings = [ZeroTrustService.__get_enriched_finding(f) for f in findings]
+        all_finding_details = ZeroTrustService._parse_finding_details_for_ui()
+        return enriched_findings
+
+
+    @staticmethod
+    def _parse_finding_details_for_ui() -> List[FindingDetails]:
+        """
+        We don't need to return all events to UI, we only display N first and N last events.
+        This code returns a list of FindingDetails with ONLY the events which are relevant to UI.
+        """
         pipeline = [{'$addFields': {'oldest_events': {'$slice': ['$events', EVENT_FETCH_CNT]},
-                                    'latest_events': {'$slice': ['$events', -1*EVENT_FETCH_CNT]},
+                                    'latest_events': {'$slice': ['$events', -1 * EVENT_FETCH_CNT]},
                                     'event_count': {'$size': '$events'}}},
                     {'$unset': ['events']}]
-        all_findings = list(Finding.objects.aggregate(*pipeline))
+        all_details = list(FindingDetails.objects.aggregate(*pipeline))
-        for finding in all_findings:
+        for details in all_details:
-            finding['latest_events'] = ZeroTrustService._get_events_without_overlap(finding['event_count'],
+            details['latest_events'] = ZeroTrustService._get_events_without_overlap(details['event_count'],
-                                                                                    finding['latest_events'])
+                                                                                    details['latest_events'])
-
-        enriched_findings = [ZeroTrustService.__get_enriched_finding(f) for f in all_findings]
-        return enriched_findings
 
     @staticmethod
     def _get_events_without_overlap(event_count: int, events: List[object]) -> List[object]:
@@ -140,9 +155,6 @@ class ZeroTrustService(object):
             'test_key': finding['test'],
             'pillars': test_info[zero_trust_consts.PILLARS_KEY],
             'status': finding['status'],
-            'latest_events': finding['latest_events'],
-            'oldest_events': finding['oldest_events'],
-            'event_count': finding['event_count']
         }
         return enriched_finding
 
@@ -169,7 +181,7 @@ class ZeroTrustService(object):
 
     @staticmethod
     def __get_status_of_single_pillar(pillar):
-        all_findings = Finding.objects().exclude('events')
+        all_findings = Finding.objects()
         grade = ZeroTrustService.__get_pillar_grade(pillar, all_findings)
         for status in zero_trust_consts.ORDERED_TEST_STATUSES:
             if grade[status] > 0:

monkey/monkey_island/cc/services/telemetry/processing/processing.py

@@ -13,6 +13,8 @@ from monkey_island.cc.services.telemetry.processing.system_info import \
     process_system_info_telemetry
 from monkey_island.cc.services.telemetry.processing.tunnel import \
     process_tunnel_telemetry
+from monkey_island.cc.services.telemetry.processing.scoutsuite import \
+    process_scoutsuite_telemetry
 
 logger = logging.getLogger(__name__)
 
@@ -24,6 +26,7 @@ TELEMETRY_CATEGORY_TO_PROCESSING_FUNC = \
         TelemCategoryEnum.SCAN: process_scan_telemetry,
         TelemCategoryEnum.SYSTEM_INFO: process_system_info_telemetry,
         TelemCategoryEnum.POST_BREACH: process_post_breach_telemetry,
+        TelemCategoryEnum.SCOUTSUITE: process_scoutsuite_telemetry,
         # `lambda *args, **kwargs: None` is a no-op.
         TelemCategoryEnum.TRACE: lambda *args, **kwargs: None,
         TelemCategoryEnum.ATTACK: lambda *args, **kwargs: None,

monkey/monkey_island/cc/services/telemetry/processing/scoutsuite.py

@@ -1,7 +1,11 @@
+import json
+
 from monkey_island.cc.database import mongo
 
 
 def process_scoutsuite_telemetry(telemetry_json):
+    # Encode data to json, because mongo can't save it as document (invalid document keys)
+    telemetry_json['data'] = json.dumps(telemetry_json['data'])
     update_data(telemetry_json)