From 9952f691981e3bdb0df867f8aecc042f1fb0aa0d Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Mon, 7 Sep 2020 13:36:18 +0300 Subject: [PATCH] Refactoring ZT findings --- .../cc/models/zero_trust/aggregate_finding.py | 22 +++++++-- .../cc/models/zero_trust/finding.py | 17 +++---- .../cc/models/zero_trust/finding_details.py | 27 +++++++++++ .../models/zero_trust/scoutsuite_finding.py | 17 +++++++ .../cc/resources/reporting/report.py | 2 +- .../monkey_island/cc/resources/telemetry.py | 1 + .../services/reporting/zero_trust_service.py | 46 ++++++++++++------- .../telemetry/processing/processing.py | 3 ++ .../telemetry/processing/scoutsuite.py | 4 ++ 9 files changed, 107 insertions(+), 32 deletions(-) create mode 100644 monkey/monkey_island/cc/models/zero_trust/finding_details.py create mode 100644 monkey/monkey_island/cc/models/zero_trust/scoutsuite_finding.py diff --git a/monkey/monkey_island/cc/models/zero_trust/aggregate_finding.py b/monkey/monkey_island/cc/models/zero_trust/aggregate_finding.py index c5684abc0..af7350830 100644 --- a/monkey/monkey_island/cc/models/zero_trust/aggregate_finding.py +++ b/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] - orig_finding.add_events(events) + AggregateFinding.add_events(existing_findings[0], 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): diff --git a/monkey/monkey_island/cc/models/zero_trust/finding.py b/monkey/monkey_island/cc/models/zero_trust/finding.py index d07cb9a45..b6bbf900d 100644 --- a/monkey/monkey_island/cc/models/zero_trust/finding.py +++ b/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): - finding = Finding( - test=test, - status=status, - events=events) + def save_finding(test: str, status: str, detail_ref): + finding = Finding(test=test, + status=status, + details=detail_ref) finding.save() return finding - - def add_events(self, events: List) -> None: - self.update(push_all__events=events) diff --git a/monkey/monkey_island/cc/models/zero_trust/finding_details.py b/monkey/monkey_island/cc/models/zero_trust/finding_details.py new file mode 100644 index 000000000..260442781 --- /dev/null +++ b/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) diff --git a/monkey/monkey_island/cc/models/zero_trust/scoutsuite_finding.py b/monkey/monkey_island/cc/models/zero_trust/scoutsuite_finding.py new file mode 100644 index 000000000..f8d0f5042 --- /dev/null +++ b/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 diff --git a/monkey/monkey_island/cc/resources/reporting/report.py b/monkey/monkey_island/cc/resources/reporting/report.py index f196fdfb6..5c25d1ff6 100644 --- a/monkey/monkey_island/cc/resources/reporting/report.py +++ b/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) diff --git a/monkey/monkey_island/cc/resources/telemetry.py b/monkey/monkey_island/cc/resources/telemetry.py index dd622b140..2686f08fd 100644 --- a/monkey/monkey_island/cc/resources/telemetry.py +++ b/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} diff --git a/monkey/monkey_island/cc/services/reporting/zero_trust_service.py b/monkey/monkey_island/cc/services/reporting/zero_trust_service.py index a70f4f8b4..bd3f12e16 100644 --- a/monkey/monkey_island/cc/services/reporting/zero_trust_service.py +++ b/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)) - for finding in all_findings: - finding['latest_events'] = ZeroTrustService._get_events_without_overlap(finding['event_count'], - finding['latest_events']) - - enriched_findings = [ZeroTrustService.__get_enriched_finding(f) for f in all_findings] - return enriched_findings + all_details = list(FindingDetails.objects.aggregate(*pipeline)) + for details in all_details: + details['latest_events'] = ZeroTrustService._get_events_without_overlap(details['event_count'], + details['latest_events']) @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: diff --git a/monkey/monkey_island/cc/services/telemetry/processing/processing.py b/monkey/monkey_island/cc/services/telemetry/processing/processing.py index 960a01517..0038273b1 100644 --- a/monkey/monkey_island/cc/services/telemetry/processing/processing.py +++ b/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, diff --git a/monkey/monkey_island/cc/services/telemetry/processing/scoutsuite.py b/monkey/monkey_island/cc/services/telemetry/processing/scoutsuite.py index ae63fe508..d0e25aebc 100644 --- a/monkey/monkey_island/cc/services/telemetry/processing/scoutsuite.py +++ b/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)