From ee6b122f0181fa4986031c56b2d86a7833044037 Mon Sep 17 00:00:00 2001 From: VakarisZ Date: Tue, 12 May 2020 15:45:51 +0300 Subject: [PATCH] Minor improvements --- .../monkey_island/cc/models/zero_trust/aggregate_finding.py | 3 +-- monkey/monkey_island/cc/models/zero_trust/finding.py | 4 ++++ .../cc/services/reporting/test_zero_trust_service.py | 6 +++--- .../cc/services/reporting/zero_trust_service.py | 6 +++--- 4 files changed, 11 insertions(+), 8 deletions(-) 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 ff3c4e4d9..c3817313f 100644 --- a/monkey/monkey_island/cc/models/zero_trust/aggregate_finding.py +++ b/monkey/monkey_island/cc/models/zero_trust/aggregate_finding.py @@ -20,8 +20,7 @@ class AggregateFinding(Finding): else: # Now we know for sure this is the only one orig_finding = existing_findings[0] - orig_finding.update(push_all__events=events) - orig_finding.save() + orig_finding.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 ae1114655..2f3261ec4 100644 --- a/monkey/monkey_island/cc/models/zero_trust/finding.py +++ b/monkey/monkey_island/cc/models/zero_trust/finding.py @@ -2,6 +2,7 @@ """ Define a Document Schema for Zero Trust findings. """ +from typing import List from mongoengine import Document, StringField, EmbeddedDocumentListField @@ -54,3 +55,6 @@ class Finding(Document): finding.save() return finding + + def add_events(self, events: List) -> None: + self.update(push_all__events=events) diff --git a/monkey/monkey_island/cc/services/reporting/test_zero_trust_service.py b/monkey/monkey_island/cc/services/reporting/test_zero_trust_service.py index 403967d8f..e40af29f4 100644 --- a/monkey/monkey_island/cc/services/reporting/test_zero_trust_service.py +++ b/monkey/monkey_island/cc/services/reporting/test_zero_trust_service.py @@ -319,9 +319,9 @@ class TestZeroTrustService(IslandTestCase): def test_get_events_without_overlap(self): monkey_island.cc.services.reporting.zero_trust_service.EVENT_FETCH_CNT = 5 - self.assertListEqual([], ZeroTrustService._ZeroTrustService__get_events_without_overlap(5, [1, 2, 3])) - self.assertListEqual([3], ZeroTrustService._ZeroTrustService__get_events_without_overlap(6, [1, 2, 3])) - self.assertListEqual([1, 2, 3, 4, 5], ZeroTrustService._ZeroTrustService__get_events_without_overlap(10, [1, 2, 3, 4, 5])) + self.assertListEqual([], ZeroTrustService._get_events_without_overlap(5, [1, 2, 3])) + self.assertListEqual([3], ZeroTrustService._get_events_without_overlap(6, [1, 2, 3])) + self.assertListEqual([1, 2, 3, 4, 5], ZeroTrustService._get_events_without_overlap(10, [1, 2, 3, 4, 5])) def compare_lists_no_order(s, t): 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 821d2104a..ee8fdd8bb 100644 --- a/monkey/monkey_island/cc/services/reporting/zero_trust_service.py +++ b/monkey/monkey_island/cc/services/reporting/zero_trust_service.py @@ -115,14 +115,14 @@ class ZeroTrustService(object): {'$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']) + 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 @staticmethod - def __get_events_without_overlap(event_count: int, events: List[object]) -> List[object]: + 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