Added aggregate finding

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()

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)

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)