Minor improvements

This commit is contained in:
VakarisZ 2020-05-12 15:45:51 +03:00
parent 2debe98052
commit ee6b122f01
4 changed files with 11 additions and 8 deletions

View File

@ -20,8 +20,7 @@ class AggregateFinding(Finding):
else: else:
# Now we know for sure this is the only one # Now we know for sure this is the only one
orig_finding = existing_findings[0] orig_finding = existing_findings[0]
orig_finding.update(push_all__events=events) orig_finding.add_events(events)
orig_finding.save()
def add_malicious_activity_to_timeline(events): def add_malicious_activity_to_timeline(events):

View File

@ -2,6 +2,7 @@
""" """
Define a Document Schema for Zero Trust findings. Define a Document Schema for Zero Trust findings.
""" """
from typing import List
from mongoengine import Document, StringField, EmbeddedDocumentListField from mongoengine import Document, StringField, EmbeddedDocumentListField
@ -54,3 +55,6 @@ class Finding(Document):
finding.save() finding.save()
return finding return finding
def add_events(self, events: List) -> None:
self.update(push_all__events=events)

View File

@ -319,9 +319,9 @@ class TestZeroTrustService(IslandTestCase):
def test_get_events_without_overlap(self): def test_get_events_without_overlap(self):
monkey_island.cc.services.reporting.zero_trust_service.EVENT_FETCH_CNT = 5 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([], ZeroTrustService._get_events_without_overlap(5, [1, 2, 3]))
self.assertListEqual([3], ZeroTrustService._ZeroTrustService__get_events_without_overlap(6, [1, 2, 3])) self.assertListEqual([3], 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([1, 2, 3, 4, 5], ZeroTrustService._get_events_without_overlap(10, [1, 2, 3, 4, 5]))
def compare_lists_no_order(s, t): def compare_lists_no_order(s, t):

View File

@ -115,14 +115,14 @@ class ZeroTrustService(object):
{'$unset': ['events']}] {'$unset': ['events']}]
all_findings = list(Finding.objects.aggregate(*pipeline)) all_findings = list(Finding.objects.aggregate(*pipeline))
for finding in all_findings: for finding in all_findings:
finding['latest_events'] = ZeroTrustService.__get_events_without_overlap(finding['event_count'], finding['latest_events'] = ZeroTrustService._get_events_without_overlap(finding['event_count'],
finding['latest_events']) finding['latest_events'])
enriched_findings = [ZeroTrustService.__get_enriched_finding(f) for f in all_findings] enriched_findings = [ZeroTrustService.__get_enriched_finding(f) for f in all_findings]
return enriched_findings return enriched_findings
@staticmethod @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 overlap_count = event_count - EVENT_FETCH_CNT
if overlap_count >= EVENT_FETCH_CNT: if overlap_count >= EVENT_FETCH_CNT:
return events return events