From 1d5a4d20cee4a86713540819307b4c77174cd5af Mon Sep 17 00:00:00 2001 From: Shay Nehmad Date: Sun, 1 Sep 2019 11:29:04 +0300 Subject: [PATCH] Added aggregate finding --- .../cc/models/zero_trust/aggregate_finding.py | 23 ++++++++ .../cc/models/zero_trust/finding.py | 4 ++ .../zero_trust/test_aggregate_finding.py | 53 +++++++++++++++++++ 3 files changed, 80 insertions(+) create mode 100644 monkey/monkey_island/cc/models/zero_trust/aggregate_finding.py create mode 100644 monkey/monkey_island/cc/models/zero_trust/test_aggregate_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 new file mode 100644 index 000000000..613b9a4a2 --- /dev/null +++ b/monkey/monkey_island/cc/models/zero_trust/aggregate_finding.py @@ -0,0 +1,23 @@ +from monkey_island.cc.models.zero_trust.finding import Finding + + +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: + Finding.save_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) + orig_finding.save() diff --git a/monkey/monkey_island/cc/models/zero_trust/finding.py b/monkey/monkey_island/cc/models/zero_trust/finding.py index 4027690c8..441d22e3a 100644 --- a/monkey/monkey_island/cc/models/zero_trust/finding.py +++ b/monkey/monkey_island/cc/models/zero_trust/finding.py @@ -54,3 +54,7 @@ class Finding(Document): finding.save() return finding + + def add_events(self, events): + # type: (list) -> None + self.events.extend(events) 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 new file mode 100644 index 000000000..b32e8ad53 --- /dev/null +++ b/monkey/monkey_island/cc/models/zero_trust/test_aggregate_finding.py @@ -0,0 +1,53 @@ +from common.data.zero_trust_consts import * +from monkey_island.cc.models.zero_trust.aggregate_finding import AggregateFinding +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 + + +class TestAggregateFinding(IslandTestCase): + def test_create_or_add_to_existing(self): + self.fail_if_not_testing_env() + self.clean_finding_db() + + test = TEST_MALICIOUS_ACTIVITY_TIMELINE + status = STATUS_INCONCLUSIVE + events = [Event.create_event("t", "t", EVENT_TYPE_ISLAND)] + self.assertEquals(len(Finding.objects(test=test, status=status)), 0) + + AggregateFinding.create_or_add_to_existing(test, status, events) + + self.assertEquals(len(Finding.objects(test=test, status=status)), 1) + self.assertEquals(len(Finding.objects(test=test, status=status)[0].events), 1) + + AggregateFinding.create_or_add_to_existing(test, status, events) + + self.assertEquals(len(Finding.objects(test=test, status=status)), 1) + self.assertEquals(len(Finding.objects(test=test, status=status)[0].events), 2) + + def test_create_or_add_to_existing_2_tests_already_exist(self): + self.fail_if_not_testing_env() + self.clean_finding_db() + + test = TEST_MALICIOUS_ACTIVITY_TIMELINE + status = STATUS_INCONCLUSIVE + event = Event.create_event("t", "t", EVENT_TYPE_ISLAND) + events = [event] + self.assertEquals(len(Finding.objects(test=test, status=status)), 0) + + Finding.save_finding(test, status, events) + + self.assertEquals(len(Finding.objects(test=test, status=status)), 1) + self.assertEquals(len(Finding.objects(test=test, status=status)[0].events), 1) + + AggregateFinding.create_or_add_to_existing(test, status, events) + + self.assertEquals(len(Finding.objects(test=test, status=status)), 1) + self.assertEquals(len(Finding.objects(test=test, status=status)[0].events), 2) + + Finding.save_finding(test, status, events) + + self.assertEquals(len(Finding.objects(test=test, status=status)), 2) + + with self.assertRaises(AssertionError): + AggregateFinding.create_or_add_to_existing(test, status, events)