ZeroTrust performance fixes

envs/monkey_zoo/blackbox/test_blackbox.py

@@ -10,7 +10,10 @@
 from envs.monkey_zoo.blackbox.log_handlers.test_logs_handler import TestLogsHandler
 from envs.monkey_zoo.blackbox.tests.exploitation import ExploitationTest
 from envs.monkey_zoo.blackbox.tests.performance.map_generation import MapGenerationTest
+from envs.monkey_zoo.blackbox.tests.performance.map_generation_from_telemetries import MapGenerationFromTelemetryTest
 from envs.monkey_zoo.blackbox.tests.performance.report_generation import ReportGenerationTest
+from envs.monkey_zoo.blackbox.tests.performance.report_generation_from_telemetries import \
+    ReportGenerationFromTelemetryTest
 from envs.monkey_zoo.blackbox.tests.performance.telemetry_performance_test import TelemetryPerformanceTest
 from envs.monkey_zoo.blackbox.utils import gcp_machine_handlers
 
@@ -146,5 +149,11 @@ def test_map_generation_performance(self, island_client):
                                                 "PERFORMANCE.conf",
                                                 timeout_in_seconds=10*60)
 
+    def test_report_generation_from_fake_telemetries(self, island_client):
+        ReportGenerationFromTelemetryTest(island_client).run()
+
+    def test_map_generation_from_fake_telemetries(self, island_client):
+        MapGenerationFromTelemetryTest(island_client).run()
+
     def test_telem_performance(self, island_client):
         TelemetryPerformanceTest(island_client).test_telemetry_performance()

envs/monkey_zoo/blackbox/tests/performance/map_generation_from_telemetries.py

@@ -0,0 +1,31 @@
+from datetime import timedelta
+
+from envs.monkey_zoo.blackbox.tests.performance.performance_test import PerformanceTest
+from envs.monkey_zoo.blackbox.tests.performance.performance_test_config import PerformanceTestConfig
+from envs.monkey_zoo.blackbox.tests.performance.telemetry_performance_test_workflow import \
+    TelemetryPerformanceTestWorkflow
+
+MAX_ALLOWED_SINGLE_PAGE_TIME = timedelta(seconds=2)
+MAX_ALLOWED_TOTAL_TIME = timedelta(seconds=5)
+
+MAP_RESOURCES = [
+    "api/netmap",
+]
+
+
+class MapGenerationFromTelemetryTest(PerformanceTest):
+
+    TEST_NAME = "Map generation from fake telemetries test"
+
+    def __init__(self, island_client, break_on_timeout=False):
+        self.island_client = island_client
+        performance_config = PerformanceTestConfig(max_allowed_single_page_time=MAX_ALLOWED_SINGLE_PAGE_TIME,
+                                                   max_allowed_total_time=MAX_ALLOWED_TOTAL_TIME,
+                                                   endpoints_to_test=MAP_RESOURCES,
+                                                   break_on_timeout=break_on_timeout)
+        self.performance_test_workflow = TelemetryPerformanceTestWorkflow(MapGenerationFromTelemetryTest.TEST_NAME,
+                                                                          self.island_client,
+                                                                          performance_config)
+
+    def run(self):
+        self.performance_test_workflow.run()

envs/monkey_zoo/blackbox/tests/performance/performance_test_workflow.py

@@ -25,6 +25,8 @@ def run(self):
             self.exploitation_test.wait_for_monkey_process_to_finish()
         performance_test = EndpointPerformanceTest(self.name, self.performance_config, self.island_client)
         try:
+            if not self.island_client.is_all_monkeys_dead():
+                raise RuntimeError("Can't test report times since not all Monkeys have died.")
             assert performance_test.run()
         finally:
             self.exploitation_test.parse_logs()

envs/monkey_zoo/blackbox/tests/performance/report_generation_from_telemetries.py

@@ -0,0 +1,35 @@
+from datetime import timedelta
+
+from envs.monkey_zoo.blackbox.tests.performance.performance_test import PerformanceTest
+from envs.monkey_zoo.blackbox.tests.performance.performance_test_config import PerformanceTestConfig
+from envs.monkey_zoo.blackbox.tests.performance.telemetry_performance_test_workflow import \
+    TelemetryPerformanceTestWorkflow
+
+MAX_ALLOWED_SINGLE_PAGE_TIME = timedelta(seconds=2)
+MAX_ALLOWED_TOTAL_TIME = timedelta(seconds=5)
+
+REPORT_RESOURCES = [
+    "api/report/security",
+    "api/attack/report",
+    "api/report/zero_trust/findings",
+    "api/report/zero_trust/principles",
+    "api/report/zero_trust/pillars"
+]
+
+
+class ReportGenerationFromTelemetryTest(PerformanceTest):
+
+    TEST_NAME = "Map generation from fake telemetries test"
+
+    def __init__(self, island_client, break_on_timeout=False):
+        self.island_client = island_client
+        performance_config = PerformanceTestConfig(max_allowed_single_page_time=MAX_ALLOWED_SINGLE_PAGE_TIME,
+                                                   max_allowed_total_time=MAX_ALLOWED_TOTAL_TIME,
+                                                   endpoints_to_test=REPORT_RESOURCES,
+                                                   break_on_timeout=break_on_timeout)
+        self.performance_test_workflow = TelemetryPerformanceTestWorkflow(ReportGenerationFromTelemetryTest.TEST_NAME,
+                                                                          self.island_client,
+                                                                          performance_config)
+
+    def run(self):
+        self.performance_test_workflow.run()

envs/monkey_zoo/blackbox/tests/performance/telemetry_performance_test_workflow.py

@@ -0,0 +1,20 @@
+from envs.monkey_zoo.blackbox.tests.basic_test import BasicTest
+from envs.monkey_zoo.blackbox.tests.performance.endpoint_performance_test import EndpointPerformanceTest
+from envs.monkey_zoo.blackbox.tests.performance.performance_test_config import PerformanceTestConfig
+from envs.monkey_zoo.blackbox.tests.performance.telemetry_performance_test import TelemetryPerformanceTest
+
+
+class TelemetryPerformanceTestWorkflow(BasicTest):
+
+    def __init__(self, name, island_client, performance_config: PerformanceTestConfig):
+        self.name = name
+        self.island_client = island_client
+        self.performance_config = performance_config
+
+    def run(self):
+        try:
+            TelemetryPerformanceTest(island_client=self.island_client).test_telemetry_performance()
+            performance_test = EndpointPerformanceTest(self.name, self.performance_config, self.island_client)
+            assert performance_test.run()
+        finally:
+            self.island_client.reset_env()

monkey/monkey_island/cc/app.py

@@ -32,6 +32,7 @@
 from monkey_island.cc.resources.attack.attack_config import AttackConfiguration
 from monkey_island.cc.resources.attack.attack_report import AttackReport
 from monkey_island.cc.resources.bootloader import Bootloader
+from monkey_island.cc.resources.zero_trust.finding_event import ZeroTrustFindingEvent
 from monkey_island.cc.services.database import Database
 from monkey_island.cc.services.remote_run_aws import RemoteRunAwsService
 from monkey_island.cc.services.representations import output_json
@@ -107,6 +108,7 @@ def init_api_resources(api):
         Report,
         '/api/report/<string:report_type>',
         '/api/report/<string:report_type>/<string:report_data>')
+    api.add_resource(ZeroTrustFindingEvent, '/api/zero-trust/finding-event/<string:finding_id>')
 
     api.add_resource(TelemetryFeed, '/api/telemetry-feed', '/api/telemetry-feed/')
     api.add_resource(Log, '/api/log', '/api/log/')

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

@@ -12,15 +12,15 @@ def create_or_add_to_existing(test, status, events):
         :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)
+        existing_findings = Finding.objects(test=test, status=status).exclude('events')
         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.update(push_all__events=events)
             orig_finding.save()
 
 

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

@@ -23,7 +23,9 @@ class Event(EmbeddedDocument):
 
     # LOGIC
     @staticmethod
-    def create_event(title, message, event_type, timestamp=datetime.now()):
+    def create_event(title, message, event_type, timestamp=None):
+        if not timestamp:
+            timestamp = datetime.now()
         event = Event(
             timestamp=timestamp,
             title=title,

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

@@ -54,7 +54,3 @@ def save_finding(test, status, events):
         finding.save()
 
         return finding
-
-    def add_events(self, events):
-        # type: (list) -> None
-        self.events.extend(events)

monkey/monkey_island/cc/resources/zero_trust/finding_event.py

@@ -0,0 +1,12 @@
+import flask_restful
+import json
+
+from monkey_island.cc.auth import jwt_required
+from monkey_island.cc.services.reporting.zero_trust_service import ZeroTrustService
+
+
+class ZeroTrustFindingEvent(flask_restful.Resource):
+
+    @jwt_required()
+    def get(self, finding_id: str):
+        return {'events_json': json.dumps(ZeroTrustService.get_events_by_finding(finding_id), default=str)}

monkey/monkey_island/cc/services/reporting/test_zero_trust_service.py

@@ -1,6 +1,7 @@
 import common.data.zero_trust_consts as zero_trust_consts
 from monkey_island.cc.models.zero_trust.finding import Finding
 from monkey_island.cc.services.reporting.zero_trust_service import ZeroTrustService
+import monkey_island.cc.services.reporting.zero_trust_service
 from monkey_island.cc.testing.IslandTestCase import IslandTestCase
 
 EXPECTED_DICT = {
@@ -316,6 +317,12 @@ def test_get_pillars_to_statuses(self):
 
         self.assertEqual(ZeroTrustService.get_pillars_to_statuses(), expected)
 
+    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]))
+
 
 def compare_lists_no_order(s, t):
     t = list(t)  # make a mutable copy

monkey/monkey_island/cc/services/reporting/zero_trust_service.py

@@ -1,21 +1,26 @@
-import json
+from typing import List
 
 import common.data.zero_trust_consts as zero_trust_consts
+from bson.objectid import ObjectId
 
 from monkey_island.cc.models.zero_trust.finding import Finding
 
+# 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 ZeroTrustService(object):
     @staticmethod
     def get_pillars_grades():
         pillars_grades = []
+        all_findings = Finding.objects().exclude('events')
         for pillar in zero_trust_consts.PILLARS:
-            pillars_grades.append(ZeroTrustService.__get_pillar_grade(pillar))
+            pillars_grades.append(ZeroTrustService.__get_pillar_grade(pillar, all_findings))
         return pillars_grades
 
     @staticmethod
-    def __get_pillar_grade(pillar):
-        all_findings = Finding.objects()
+    def __get_pillar_grade(pillar, all_findings):
         pillar_grade = {
             "pillar": pillar,
             zero_trust_consts.STATUS_FAILED: 0,
@@ -65,7 +70,7 @@ def __get_principle_status(principle_tests):
         worst_status = zero_trust_consts.STATUS_UNEXECUTED
         all_statuses = set()
         for test in principle_tests:
-            all_statuses |= set(Finding.objects(test=test).distinct("status"))
+            all_statuses |= set(Finding.objects(test=test).exclude('events').distinct("status"))
 
         for status in all_statuses:
             if zero_trust_consts.ORDERED_TEST_STATUSES.index(status) \
@@ -78,7 +83,7 @@ def __get_principle_status(principle_tests):
     def __get_tests_status(principle_tests):
         results = []
         for test in principle_tests:
-            test_findings = Finding.objects(test=test)
+            test_findings = Finding.objects(test=test).exclude('events')
             results.append(
                 {
                     "test": zero_trust_consts.TESTS_MAP[test][zero_trust_consts.TEST_EXPLANATION_KEY],
@@ -104,26 +109,43 @@ def __get_lcd_worst_status_for_test(all_findings_for_test):
 
     @staticmethod
     def get_all_findings():
-        all_findings = Finding.objects()
+        pipeline = [{'$addFields': {'oldest_events': {'$slice': ['$events', EVENT_FETCH_CNT]},
+                                    'latest_events': {'$slice': ['$events', -1*EVENT_FETCH_CNT]},
+                                    'event_count': {'$size': '$events'}}},
+                    {'$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'])
+
         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]:
+        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:]
+
     @staticmethod
     def __get_enriched_finding(finding):
-        test_info = zero_trust_consts.TESTS_MAP[finding.test]
+        test_info = zero_trust_consts.TESTS_MAP[finding['test']]
         enriched_finding = {
-            "test": test_info[zero_trust_consts.FINDING_EXPLANATION_BY_STATUS_KEY][finding.status],
-            "test_key": finding.test,
-            "pillars": test_info[zero_trust_consts.PILLARS_KEY],
-            "status": finding.status,
-            "events": ZeroTrustService.__get_events_as_dict(finding.events)
+            'finding_id': str(finding['_id']),
+            'test': test_info[zero_trust_consts.FINDING_EXPLANATION_BY_STATUS_KEY][finding['status']],
+            'test_key': finding['test'],
+            'pillars': test_info[zero_trust_consts.PILLARS_KEY],
+            'status': finding['status'],
+            'latest_events': finding['latest_events'],
+            'oldest_events': finding['oldest_events'],
+            'event_count': finding['event_count']
         }
         return enriched_finding
 
-    @staticmethod
-    def __get_events_as_dict(events):
-        return [json.loads(event.to_json()) for event in events]
-
     @staticmethod
     def get_statuses_to_pillars():
         results = {
@@ -147,8 +169,17 @@ def get_pillars_to_statuses():
 
     @staticmethod
     def __get_status_of_single_pillar(pillar):
-        grade = ZeroTrustService.__get_pillar_grade(pillar)
+        all_findings = Finding.objects().exclude('events')
+        grade = ZeroTrustService.__get_pillar_grade(pillar, all_findings)
         for status in zero_trust_consts.ORDERED_TEST_STATUSES:
             if grade[status] > 0:
                 return status
         return zero_trust_consts.STATUS_UNEXECUTED
+
+    @staticmethod
+    def get_events_by_finding(finding_id: str) -> List[object]:
+        pipeline = [{'$match': {'_id': ObjectId(finding_id)}},
+                    {'$unwind': '$events'},
+                    {'$project': {'events': '$events'}},
+                    {'$replaceRoot': {'newRoot': '$events'}}]
+        return list(Finding.objects.aggregate(*pipeline))

monkey/monkey_island/cc/testing/profiler_decorator.py

@@ -0,0 +1,32 @@
+from cProfile import Profile
+import os
+import pstats
+
+PROFILER_LOG_DIR = "./profiler_logs/"
+
+
+def profile(sort_args=['cumulative'], print_args=[100]):
+
+    def decorator(fn):
+        def inner(*args, **kwargs):
+            result = None
+            try:
+                profiler = Profile()
+                result = profiler.runcall(fn, *args, **kwargs)
+            finally:
+                try:
+                    os.mkdir(PROFILER_LOG_DIR)
+                except os.error:
+                    pass
+                filename = PROFILER_LOG_DIR + _get_filename_for_function(fn)
+                with open(filename, 'w') as stream:
+                    stats = pstats.Stats(profiler, stream=stream)
+                    stats.strip_dirs().sort_stats(*sort_args).print_stats(*print_args)
+            return result
+        return inner
+    return decorator
+
+
+def _get_filename_for_function(fn):
+    function_name = fn.__module__ + "." + fn.__name__
+    return function_name.replace(".", "_")

monkey/monkey_island/cc/ui/src/components/report-components/zerotrust/EventsButton.js

@@ -24,7 +24,12 @@ export default class EventsButton extends Component {
 
   render() {
     return <Fragment>
-      <EventsModal events={this.props.events} showEvents={this.state.isShow} hideCallback={this.hide}
+      <EventsModal finding_id={this.props.finding_id}
+                   latest_events={this.props.latest_events}
+                   oldest_events={this.props.oldest_events}
+                   event_count={this.props.event_count}
+                   showEvents={this.state.isShow}
+                   hideCallback={this.hide}
                    exportFilename={this.props.exportFilename}/>
       <div className="text-center" style={{'display': 'grid'}}>
         <Button className="btn btn-primary btn-lg" onClick={this.show}>
@@ -35,12 +40,14 @@ export default class EventsButton extends Component {
   }
 
   createEventsAmountBadge() {
-    const eventsAmountBadgeContent = this.props.events.length > 9 ? '9+' : this.props.events.length;
+    const eventsAmountBadgeContent = this.props.event_count > 9 ? '9+' : this.props.event_count;
     return <Badge>{eventsAmountBadgeContent}</Badge>;
   }
 }
 
 EventsButton.propTypes = {
-  events: PropTypes.array,
+  latest_events: PropTypes.array,
+  oldest_events: PropTypes.array,
+  event_count: PropTypes.number,
   exportFilename: PropTypes.string
 };