From 3490be1d8f499be429ad3ce2a5cdf82cbede1060 Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Tue, 8 Sep 2020 12:39:55 +0300 Subject: [PATCH] Re-structured ZT files and separated class responsibilities better, also further refactor towards ZT findings being extendable with different types of details. --- .../cc/models/zero_trust/aggregate_finding.py | 45 ------------ .../cc/models/zero_trust/finding.py | 7 +- .../zero_trust/monkey_finding_details.py | 18 +++++ ...tails.py => scoutsuite_finding_details.py} | 10 +-- .../zero_trust/test_aggregate_finding.py | 11 ++- .../cc/resources/reporting/report.py | 2 +- .../cc/resources/zero_trust/finding_event.py | 5 +- .../__init__.py | 0 .../cc/services/zero_trust/events_service.py | 34 ++++++++++ .../cc/services/zero_trust/finding_service.py | 23 +++++++ .../zero_trust/monkey_finding_service.py | 68 +++++++++++++++++++ .../zero_trust/scoutsuite_finding_service.py | 3 + .../test_zero_trust_service.py | 4 +- .../zero_trust_service.py | 64 +---------------- 14 files changed, 162 insertions(+), 132 deletions(-) delete mode 100644 monkey/monkey_island/cc/models/zero_trust/aggregate_finding.py create mode 100644 monkey/monkey_island/cc/models/zero_trust/monkey_finding_details.py rename monkey/monkey_island/cc/models/zero_trust/{finding_details.py => scoutsuite_finding_details.py} (63%) rename monkey/monkey_island/cc/services/{telemetry/zero_trust_tests => zero_trust}/__init__.py (100%) create mode 100644 monkey/monkey_island/cc/services/zero_trust/events_service.py create mode 100644 monkey/monkey_island/cc/services/zero_trust/finding_service.py create mode 100644 monkey/monkey_island/cc/services/zero_trust/monkey_finding_service.py create mode 100644 monkey/monkey_island/cc/services/zero_trust/scoutsuite_finding_service.py rename monkey/monkey_island/cc/services/{reporting => zero_trust}/test_zero_trust_service.py (99%) rename monkey/monkey_island/cc/services/{reporting => zero_trust}/zero_trust_service.py (64%) diff --git a/monkey/monkey_island/cc/models/zero_trust/aggregate_finding.py b/monkey/monkey_island/cc/models/zero_trust/aggregate_finding.py deleted file mode 100644 index af7350830..000000000 --- a/monkey/monkey_island/cc/models/zero_trust/aggregate_finding.py +++ /dev/null @@ -1,45 +0,0 @@ -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): - @staticmethod - def create_or_add_to_existing(test, status, events): - """ - Create a new finding or add the events to an existing one if it's the same (same meaning same status and same - test). - - :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) - assert (len(existing_findings) < 2), "More than one finding exists for {}:{}".format(test, status) - - if len(existing_findings) == 0: - AggregateFinding.create_new_finding(test, status, events) - else: - # Now we know for sure this is the only one - 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): - AggregateFinding.create_or_add_to_existing( - test=zero_trust_consts.TEST_MALICIOUS_ACTIVITY_TIMELINE, - status=zero_trust_consts.STATUS_VERIFY, - events=events - ) diff --git a/monkey/monkey_island/cc/models/zero_trust/finding.py b/monkey/monkey_island/cc/models/zero_trust/finding.py index b6bbf900d..8895a7fdb 100644 --- a/monkey/monkey_island/cc/models/zero_trust/finding.py +++ b/monkey/monkey_island/cc/models/zero_trust/finding.py @@ -4,13 +4,14 @@ Define a Document Schema for Zero Trust findings. """ from typing import List -from mongoengine import Document, EmbeddedDocumentListField, StringField, LazyReferenceField +from mongoengine import Document, StringField, GenericLazyReferenceField 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 +from monkey_island.cc.models.zero_trust.monkey_finding_details import MonkeyFindingDetails +from monkey_island.cc.models.zero_trust.scoutsuite_finding_details import ScoutsuiteFindingDetails class Finding(Document): @@ -34,7 +35,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) - details = LazyReferenceField(document_type=FindingDetails, required=True) + details = GenericLazyReferenceField(choices=[MonkeyFindingDetails, ScoutsuiteFindingDetails], required=True) # http://docs.mongoengine.org/guide/defining-documents.html#document-inheritance meta = {'allow_inheritance': True} diff --git a/monkey/monkey_island/cc/models/zero_trust/monkey_finding_details.py b/monkey/monkey_island/cc/models/zero_trust/monkey_finding_details.py new file mode 100644 index 000000000..029136679 --- /dev/null +++ b/monkey/monkey_island/cc/models/zero_trust/monkey_finding_details.py @@ -0,0 +1,18 @@ +from typing import List + +from mongoengine import DateTimeField, Document, StringField, EmbeddedDocumentListField + +from monkey_island.cc.models.zero_trust.event import Event + +class MonkeyFindingDetails(Document): + """ + This model represents additional information about monkey finding: + Events + """ + + # SCHEMA + events = EmbeddedDocumentListField(document_type=Event, required=False) + + # LOGIC + def add_events(self, events: List[Event]) -> 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/scoutsuite_finding_details.py similarity index 63% rename from monkey/monkey_island/cc/models/zero_trust/finding_details.py rename to monkey/monkey_island/cc/models/zero_trust/scoutsuite_finding_details.py index 260442781..ed05f6003 100644 --- a/monkey/monkey_island/cc/models/zero_trust/finding_details.py +++ b/monkey/monkey_island/cc/models/zero_trust/scoutsuite_finding_details.py @@ -1,14 +1,11 @@ -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): +class ScoutsuiteFindingDetails(Document): """ This model represents additional information about monkey finding: Events if monkey finding @@ -16,12 +13,7 @@ class FindingDetails(Document): """ # 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/test_aggregate_finding.py b/monkey/monkey_island/cc/models/zero_trust/test_aggregate_finding.py index 4b9765f70..5d042312f 100644 --- a/monkey/monkey_island/cc/models/zero_trust/test_aggregate_finding.py +++ b/monkey/monkey_island/cc/models/zero_trust/test_aggregate_finding.py @@ -4,8 +4,7 @@ import mongomock from packaging import version import common.common_consts.zero_trust_consts as zero_trust_consts -from monkey_island.cc.models.zero_trust.aggregate_finding import \ - AggregateFinding +from monkey_island.cc.services.zero_trust.monkey_finding_service import MonkeyFindingService from monkey_island.cc.models.zero_trust.event import Event from monkey_island.cc.models.zero_trust.finding import Finding from monkey_island.cc.testing.IslandTestCase import IslandTestCase @@ -24,12 +23,12 @@ class TestAggregateFinding(IslandTestCase): events = [Event.create_event("t", "t", zero_trust_consts.EVENT_TYPE_MONKEY_NETWORK)] self.assertEqual(len(Finding.objects(test=test, status=status)), 0) - AggregateFinding.create_or_add_to_existing(test, status, events) + MonkeyFindingService.create_or_add_to_existing(test, status, events) self.assertEqual(len(Finding.objects(test=test, status=status)), 1) self.assertEqual(len(Finding.objects(test=test, status=status)[0].events), 1) - AggregateFinding.create_or_add_to_existing(test, status, events) + MonkeyFindingService.create_or_add_to_existing(test, status, events) self.assertEqual(len(Finding.objects(test=test, status=status)), 1) self.assertEqual(len(Finding.objects(test=test, status=status)[0].events), 2) @@ -51,7 +50,7 @@ class TestAggregateFinding(IslandTestCase): self.assertEqual(len(Finding.objects(test=test, status=status)), 1) self.assertEqual(len(Finding.objects(test=test, status=status)[0].events), 1) - AggregateFinding.create_or_add_to_existing(test, status, events) + MonkeyFindingService.create_or_add_to_existing(test, status, events) self.assertEqual(len(Finding.objects(test=test, status=status)), 1) self.assertEqual(len(Finding.objects(test=test, status=status)[0].events), 2) @@ -61,4 +60,4 @@ class TestAggregateFinding(IslandTestCase): self.assertEqual(len(Finding.objects(test=test, status=status)), 2) with self.assertRaises(AssertionError): - AggregateFinding.create_or_add_to_existing(test, status, events) + MonkeyFindingService.create_or_add_to_existing(test, status, events) diff --git a/monkey/monkey_island/cc/resources/reporting/report.py b/monkey/monkey_island/cc/resources/reporting/report.py index 5c25d1ff6..f1b95607f 100644 --- a/monkey/monkey_island/cc/resources/reporting/report.py +++ b/monkey/monkey_island/cc/resources/reporting/report.py @@ -5,7 +5,7 @@ from flask import jsonify from monkey_island.cc.resources.auth.auth import jwt_required from monkey_island.cc.services.reporting.report import ReportService -from monkey_island.cc.services.reporting.zero_trust_service import \ +from monkey_island.cc.services.zero_trust.zero_trust_service import \ ZeroTrustService ZERO_TRUST_REPORT_TYPE = "zero_trust" diff --git a/monkey/monkey_island/cc/resources/zero_trust/finding_event.py b/monkey/monkey_island/cc/resources/zero_trust/finding_event.py index 8a1879c9c..0e6c09b11 100644 --- a/monkey/monkey_island/cc/resources/zero_trust/finding_event.py +++ b/monkey/monkey_island/cc/resources/zero_trust/finding_event.py @@ -3,12 +3,11 @@ import json import flask_restful from monkey_island.cc.resources.auth.auth import jwt_required -from monkey_island.cc.services.reporting.zero_trust_service import \ - ZeroTrustService +from monkey_island.cc.services.zero_trust.monkey_finding_service import MonkeyFindingService class ZeroTrustFindingEvent(flask_restful.Resource): @jwt_required def get(self, finding_id: str): - return {'events_json': json.dumps(ZeroTrustService.get_events_by_finding(finding_id), default=str)} + return {'events_json': json.dumps(MonkeyFindingService.get_events_by_finding(finding_id), default=str)} diff --git a/monkey/monkey_island/cc/services/telemetry/zero_trust_tests/__init__.py b/monkey/monkey_island/cc/services/zero_trust/__init__.py similarity index 100% rename from monkey/monkey_island/cc/services/telemetry/zero_trust_tests/__init__.py rename to monkey/monkey_island/cc/services/zero_trust/__init__.py diff --git a/monkey/monkey_island/cc/services/zero_trust/events_service.py b/monkey/monkey_island/cc/services/zero_trust/events_service.py new file mode 100644 index 000000000..5cccee7f3 --- /dev/null +++ b/monkey/monkey_island/cc/services/zero_trust/events_service.py @@ -0,0 +1,34 @@ +from typing import List + +from bson import ObjectId + +from monkey_island.cc.models.zero_trust.monkey_finding_details import MonkeyFindingDetails + +# How many events of a single finding to return to UI. +# 50 will return 50 latest and 50 oldest events from a finding +EVENT_FETCH_CNT = 50 + + +class EventsService: + + @staticmethod + def fetch_events_for_display(finding_id: ObjectId): + pipeline = [{'$match': {'_id': finding_id}}, + {'$addFields': {'oldest_events': {'$slice': ['$events', EVENT_FETCH_CNT]}, + 'latest_events': {'$slice': ['$events', -1 * EVENT_FETCH_CNT]}, + 'event_count': {'$size': '$events'}}}, + {'$unset': ['events']}] + details = MonkeyFindingDetails.objects.aggregate(*pipeline).next() + details['latest_events'] = EventsService._get_events_without_overlap(details['event_count'], + details['latest_events']) + return details + + @staticmethod + def _get_events_without_overlap(event_count: int, events: List[object]) -> List[object]: + overlap_count = event_count - EVENT_FETCH_CNT + if overlap_count >= EVENT_FETCH_CNT: + return events + elif overlap_count <= 0: + return [] + else: + return events[-1 * overlap_count:] diff --git a/monkey/monkey_island/cc/services/zero_trust/finding_service.py b/monkey/monkey_island/cc/services/zero_trust/finding_service.py new file mode 100644 index 000000000..2feb02cce --- /dev/null +++ b/monkey/monkey_island/cc/services/zero_trust/finding_service.py @@ -0,0 +1,23 @@ +from typing import List + +from common.common_consts import zero_trust_consts +from monkey_island.cc.models.zero_trust.finding import Finding + + +class FindingService: + + @staticmethod + def get_all_findings() -> List[Finding]: + return list(Finding.objects) + + @staticmethod + def get_enriched_finding(finding): + test_info = zero_trust_consts.TESTS_MAP[finding['test']] + enriched_finding = { + 'finding_id': str(finding['_id']), + 'test': test_info[zero_trust_consts.FINDING_EXPLANATION_BY_STATUS_KEY][finding['status']], + 'test_key': finding['test'], + 'pillars': test_info[zero_trust_consts.PILLARS_KEY], + 'status': finding['status'], + } + return enriched_finding diff --git a/monkey/monkey_island/cc/services/zero_trust/monkey_finding_service.py b/monkey/monkey_island/cc/services/zero_trust/monkey_finding_service.py new file mode 100644 index 000000000..a1e731b14 --- /dev/null +++ b/monkey/monkey_island/cc/services/zero_trust/monkey_finding_service.py @@ -0,0 +1,68 @@ +from typing import List + +from bson import ObjectId + +from common.common_consts import 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.monkey_finding_details import MonkeyFindingDetails +from monkey_island.cc.services.zero_trust.finding_service import FindingService + + +class MonkeyFindingService: + + @staticmethod + def create_or_add_to_existing(test, status, events): + """ + Create a new finding or add the events to an existing one if it's the same (same meaning same status and same + test). + + :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) + assert (len(existing_findings) < 2), "More than one finding exists for {}:{}".format(test, status) + + if len(existing_findings) == 0: + MonkeyFindingService.create_new_finding(test, status, events) + else: + # Now we know for sure this is the only one + MonkeyFindingService.add_events(existing_findings[0], events) + + @staticmethod + def create_new_finding(test: str, status: str, events: List[Event]): + details = MonkeyFindingDetails() + 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) + + @staticmethod + def get_all_monkey_findings(): + findings = FindingService.get_all_findings() + for i in range(len(findings)): + details = MonkeyFindingService.fetch_events_for_display(findings[i].details.id) + findings[i] = findings[i].to_mongo() + findings[i] = FindingService.get_enriched_finding(findings[i]) + findings[i]['details'] = details + return findings + + @staticmethod + def get_events_by_finding(finding_id: str) -> List[object]: + finding = Finding.objects.get(id=finding_id) + pipeline = [{'$match': {'_id': ObjectId(finding.details.id)}}, + {'$unwind': '$events'}, + {'$project': {'events': '$events'}}, + {'$replaceRoot': {'newRoot': '$events'}}] + return list(MonkeyFindingDetails.objects.aggregate(*pipeline)) + + @staticmethod + def add_malicious_activity_to_timeline(events): + MonkeyFindingService.create_or_add_to_existing( + test=zero_trust_consts.TEST_MALICIOUS_ACTIVITY_TIMELINE, + status=zero_trust_consts.STATUS_VERIFY, + events=events + ) diff --git a/monkey/monkey_island/cc/services/zero_trust/scoutsuite_finding_service.py b/monkey/monkey_island/cc/services/zero_trust/scoutsuite_finding_service.py new file mode 100644 index 000000000..12ab2743b --- /dev/null +++ b/monkey/monkey_island/cc/services/zero_trust/scoutsuite_finding_service.py @@ -0,0 +1,3 @@ + +class ScoutsuiteFindingService: + pass diff --git a/monkey/monkey_island/cc/services/reporting/test_zero_trust_service.py b/monkey/monkey_island/cc/services/zero_trust/test_zero_trust_service.py similarity index 99% rename from monkey/monkey_island/cc/services/reporting/test_zero_trust_service.py rename to monkey/monkey_island/cc/services/zero_trust/test_zero_trust_service.py index 874eee293..8b3d33ba2 100644 --- a/monkey/monkey_island/cc/services/reporting/test_zero_trust_service.py +++ b/monkey/monkey_island/cc/services/zero_trust/test_zero_trust_service.py @@ -1,7 +1,7 @@ import common.common_consts.zero_trust_consts as zero_trust_consts -import monkey_island.cc.services.reporting.zero_trust_service +import monkey_island.cc.services.zero_trust.zero_trust_service from monkey_island.cc.models.zero_trust.finding import Finding -from monkey_island.cc.services.reporting.zero_trust_service import \ +from monkey_island.cc.services.zero_trust.zero_trust_service import \ ZeroTrustService from monkey_island.cc.testing.IslandTestCase import IslandTestCase diff --git a/monkey/monkey_island/cc/services/reporting/zero_trust_service.py b/monkey/monkey_island/cc/services/zero_trust/zero_trust_service.py similarity index 64% rename from monkey/monkey_island/cc/services/reporting/zero_trust_service.py rename to monkey/monkey_island/cc/services/zero_trust/zero_trust_service.py index bd3f12e16..613132a54 100644 --- a/monkey/monkey_island/cc/services/reporting/zero_trust_service.py +++ b/monkey/monkey_island/cc/services/zero_trust/zero_trust_service.py @@ -5,14 +5,8 @@ from bson.objectid import ObjectId import common.common_consts.zero_trust_consts as zero_trust_consts 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 - - -class ZeroTrustService(object): +class ZeroTrustService: @staticmethod def get_pillars_grades(): pillars_grades = [] @@ -110,54 +104,6 @@ class ZeroTrustService(object): return current_worst_status - @staticmethod - 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]}, - 'event_count': {'$size': '$events'}}}, - {'$unset': ['events']}] - 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]: - overlap_count = event_count - EVENT_FETCH_CNT - if overlap_count >= EVENT_FETCH_CNT: - return events - elif overlap_count <= 0: - return [] - else: - return events[-1 * overlap_count:] - - @staticmethod - def __get_enriched_finding(finding): - test_info = zero_trust_consts.TESTS_MAP[finding['test']] - enriched_finding = { - 'finding_id': str(finding['_id']), - 'test': test_info[zero_trust_consts.FINDING_EXPLANATION_BY_STATUS_KEY][finding['status']], - 'test_key': finding['test'], - 'pillars': test_info[zero_trust_consts.PILLARS_KEY], - 'status': finding['status'], - } - return enriched_finding - @staticmethod def get_statuses_to_pillars(): results = { @@ -187,11 +133,3 @@ class ZeroTrustService(object): if grade[status] > 0: return status return zero_trust_consts.STATUS_UNEXECUTED - - @staticmethod - def get_events_by_finding(finding_id: str) -> List[object]: - pipeline = [{'$match': {'_id': ObjectId(finding_id)}}, - {'$unwind': '$events'}, - {'$project': {'events': '$events'}}, - {'$replaceRoot': {'newRoot': '$events'}}] - return list(Finding.objects.aggregate(*pipeline))