Fixed segmentation findings to use the same infrastructure as other findings.

Small segmentation finding bugfix

monkey/monkey_island/cc/models/zero_trust/monkey_finding_details.py

@@ -1,8 +1,10 @@
+from __future__ import annotations
 from typing import List
 
 from mongoengine import Document, EmbeddedDocumentListField
 
 from monkey_island.cc.models.zero_trust.event import Event
+from monkey_island.cc.models.zero_trust.subnet_pair import SubnetPair
 
 
 class MonkeyFindingDetails(Document):
@@ -13,8 +15,16 @@ class MonkeyFindingDetails(Document):
 
     # SCHEMA
     events = EmbeddedDocumentListField(document_type=Event, required=False)
+    checked_subnet_pairs = EmbeddedDocumentListField(document_type=SubnetPair, required=False)
 
     # LOGIC
-    def add_events(self, events: List[Event]) -> None:
+    def add_events(self, events: List[Event]) -> MonkeyFindingDetails:
         self.update(push_all__events=events)
         return self
+
+    def add_checked_subnet_pair(self, subnet_pair: SubnetPair) -> MonkeyFindingDetails:
+        self.update(push__checked_subnet_pairs=subnet_pair)
+        return self
+
+    def is_with_subnet_pair(self, subnet_pair: SubnetPair):
+        return subnet_pair in self.checked_subnet_pairs

monkey/monkey_island/cc/models/zero_trust/segmentation_finding.py

@@ -1,50 +0,0 @@
-from mongoengine import StringField
-
-import common.common_consts.zero_trust_consts as zero_trust_consts
-from monkey_island.cc.models.zero_trust.finding import Finding
-
-
-def need_to_overwrite_status(saved_status, new_status):
-    return (saved_status == zero_trust_consts.STATUS_PASSED) and (new_status == zero_trust_consts.STATUS_FAILED)
-
-
-class SegmentationFinding(Finding):
-    first_subnet = StringField()
-    second_subnet = StringField()
-
-    @staticmethod
-    def create_or_add_to_existing_finding(subnets, status, segmentation_event):
-        """
-        Creates a segmentation finding. If a segmentation finding with the relevant subnets already exists, adds the
-        event to the existing finding, and the "worst" status is chosen (i.e. if the existing one is "Failed" it will
-        remain so).
-
-        :param subnets: the 2 subnets of this finding.
-        :param status: STATUS_PASSED or STATUS_FAILED
-        :param segmentation_event: The specific event
-        """
-        assert len(subnets) == 2
-
-        # Sort them so A -> B and B -> A segmentation findings will be the same one.
-        subnets.sort()
-
-        existing_findings = SegmentationFinding.objects(first_subnet=subnets[0], second_subnet=subnets[1])
-
-        if len(existing_findings) == 0:
-            # No finding exists - create.
-            new_finding = SegmentationFinding(
-                first_subnet=subnets[0],
-                second_subnet=subnets[1],
-                test=zero_trust_consts.TEST_SEGMENTATION,
-                status=status,
-                events=[segmentation_event]
-            )
-            new_finding.save()
-        else:
-            # A finding exists (should be one). Add the event to it.
-            assert len(existing_findings) == 1
-            existing_finding = existing_findings[0]
-            existing_finding.events.append(segmentation_event)
-            if need_to_overwrite_status(existing_finding.status, status):
-                existing_finding.status = status
-            existing_finding.save()

monkey/monkey_island/cc/models/zero_trust/subnet_pair.py

@@ -0,0 +1,19 @@
+from typing import List
+
+from mongoengine import EmbeddedDocument, StringField
+
+
+class SubnetPair(EmbeddedDocument):
+    """
+    This model represents a pair of subnets. It is meant to hold details about cross-segmentation check between two
+    subnets.
+    """
+    # SCHEMA
+    first_subnet = StringField()
+    second_subnet = StringField()
+
+    # LOGIC
+    @staticmethod
+    def create_subnet_pair(subnets: List[str]):
+        subnets.sort()
+        return SubnetPair(first_subnet=subnets[0], second_subnet=subnets[1])

monkey/monkey_island/cc/services/telemetry/zero_trust_checks/segmentation.py

@@ -5,8 +5,8 @@ from common.network.network_range import NetworkRange
 from common.network.segmentation_utils import get_ip_if_in_subnet, get_ip_in_src_and_not_in_dst
 from monkey_island.cc.models import Monkey
 from monkey_island.cc.models.zero_trust.event import Event
-from monkey_island.cc.models.zero_trust.segmentation_finding import SegmentationFinding
 from monkey_island.cc.services.configuration.utils import get_config_network_segments_as_subnet_groups
+from monkey_island.cc.services.zero_trust.monkey_finding_service import MonkeyFindingService
 
 SEGMENTATION_DONE_EVENT_TEXT = "Monkey on {hostname} is done attempting cross-segment communications " \
                                "from `{src_seg}` segments to `{dst_seg}` segments."
@@ -26,10 +26,11 @@ def check_segmentation_violation(current_monkey, target_ip):
             target_subnet = subnet_pair[1]
             if is_segmentation_violation(current_monkey, target_ip, source_subnet, target_subnet):
                 event = get_segmentation_violation_event(current_monkey, source_subnet, target_ip, target_subnet)
-                SegmentationFinding.create_or_add_to_existing_finding(
+                MonkeyFindingService.create_or_add_to_existing(
-                    subnets=[source_subnet, target_subnet],
+                    test=zero_trust_consts.TEST_SEGMENTATION,
+                    subnet_pairs=[[source_subnet, target_subnet]],
                     status=zero_trust_consts.STATUS_FAILED,
-                    segmentation_event=event
+                    events=[event]
                 )
 
 
@@ -90,10 +91,11 @@ def create_or_add_findings_for_all_pairs(all_subnets, current_monkey):
     all_subnets_pairs_for_this_monkey = itertools.product(this_monkey_subnets, other_subnets)
 
     for subnet_pair in all_subnets_pairs_for_this_monkey:
-        SegmentationFinding.create_or_add_to_existing_finding(
+        MonkeyFindingService.create_or_add_to_existing(
-            subnets=list(subnet_pair),
+            subnet_pairs=[list(subnet_pair)],
             status=zero_trust_consts.STATUS_PASSED,
-            segmentation_event=get_segmentation_done_event(current_monkey, subnet_pair)
+            events=[get_segmentation_done_event(current_monkey, subnet_pair)],
+            test=zero_trust_consts.TEST_SEGMENTATION
         )
 
 

monkey/monkey_island/cc/services/zero_trust/events_service.py

@@ -1,36 +0,0 @@
-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 = list(MonkeyFindingDetails.objects.aggregate(*pipeline))
-        if details:
-            details = details[0]
-            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:]

monkey/monkey_island/cc/services/zero_trust/finding_service.py

@@ -2,7 +2,7 @@ from typing import List
 
 from common.common_consts import zero_trust_consts
 from monkey_island.cc.models.zero_trust.finding import Finding
-from monkey_island.cc.services.zero_trust.events_service import EventsService
+from monkey_island.cc.services.zero_trust.monkey_details_service import MonkeyDetailsService
 
 
 class FindingService:
@@ -13,7 +13,7 @@ class FindingService:
         details = []
         for i in range(len(findings)):
             if findings[i].type == zero_trust_consts.MONKEY_FINDING:
-                details = EventsService.fetch_events_for_display(findings[i].details.id)
+                details = MonkeyDetailsService.fetch_details_for_display(findings[i].details.id)
             elif findings[i].type == zero_trust_consts.SCOUTSUITE_FINDING:
                 details = findings[i].details.fetch().to_mongo()
             findings[i] = findings[i].to_mongo()

monkey/monkey_island/cc/services/zero_trust/monkey_finding_service.py

@@ -6,12 +6,13 @@ 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.models.zero_trust.subnet_pair import SubnetPair
 
 
 class MonkeyFindingService:
 
     @staticmethod
-    def create_or_add_to_existing(test, status, events):
+    def create_or_add_to_existing(test, status, events, subnet_pairs=None):
         """
         Create a new finding or add the events to an existing one if it's the same (same meaning same status and same
         test).
@@ -23,15 +24,18 @@ class MonkeyFindingService:
         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)
+            MonkeyFindingService.create_new_finding(test, status, events, subnet_pairs)
         else:
             # Now we know for sure this is the only one
             MonkeyFindingService.add_events(existing_findings[0], events)
+            if subnet_pairs:
+                MonkeyFindingService.add_subnet_pairs(existing_findings[0], subnet_pairs)
 
     @staticmethod
-    def create_new_finding(test: str, status: str, events: List[Event]):
+    def create_new_finding(test: str, status: str, events: List[Event], subnet_pairs: List[SubnetPair]):
         details = MonkeyFindingDetails()
         details.events = events
+        details.subnet_pairs = subnet_pairs
         details.save()
         Finding.save_finding(test, status, details)
 
@@ -39,6 +43,20 @@ class MonkeyFindingService:
     def add_events(finding: Finding, events: List[Event]):
         finding.details.fetch().add_events(events).save()
 
+    @staticmethod
+    def add_subnet_pairs(finding: Finding, subnet_pairs: List[List[str]]):
+        finding_details = finding.details.fetch()
+        for subnet_pair in subnet_pairs:
+            subnet_pair_document = SubnetPair.create_subnet_pair(subnet_pair)
+            if not MonkeyFindingService.is_subnet_pair_in_finding(finding, subnet_pair_document):
+                finding_details.add_checked_subnet_pair(subnet_pair_document)
+        finding_details.save()
+
+    @staticmethod
+    def is_subnet_pair_in_finding(finding: Finding, subnet_pair: SubnetPair):
+        details = finding.details.fetch()
+        return details.is_with_subnet_pair(subnet_pair)
+
     @staticmethod
     def get_events_by_finding(finding_id: str) -> List[object]:
         finding = Finding.objects.get(id=finding_id)
@@ -50,8 +68,5 @@ class MonkeyFindingService:
 
     @staticmethod
     def add_malicious_activity_to_timeline(events):
-        MonkeyFindingService.create_or_add_to_existing(
+        MonkeyFindingService.create_or_add_to_existing(test=zero_trust_consts.TEST_MALICIOUS_ACTIVITY_TIMELINE,
-            test=zero_trust_consts.TEST_MALICIOUS_ACTIVITY_TIMELINE,
+                                                       status=zero_trust_consts.STATUS_VERIFY, events=events)
-            status=zero_trust_consts.STATUS_VERIFY,
-            events=events
-        )