diff --git a/.gitignore b/.gitignore index a1012776721..960e8c67c0d 100644 --- a/.gitignore +++ b/.gitignore @@ -85,5 +85,8 @@ MonkeyZoo/* # Exported monkey telemetries /monkey/telem_sample/ +# Profiling logs +profiler_logs/ + # vim swap files *.swp diff --git a/envs/monkey_zoo/blackbox/test_blackbox.py b/envs/monkey_zoo/blackbox/test_blackbox.py index 9e3211691f4..2408d79bed2 100644 --- a/envs/monkey_zoo/blackbox/test_blackbox.py +++ b/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() diff --git a/envs/monkey_zoo/blackbox/tests/performance/map_generation_from_telemetries.py b/envs/monkey_zoo/blackbox/tests/performance/map_generation_from_telemetries.py new file mode 100644 index 00000000000..c5344d8f7f3 --- /dev/null +++ b/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() diff --git a/envs/monkey_zoo/blackbox/tests/performance/performance_test_workflow.py b/envs/monkey_zoo/blackbox/tests/performance/performance_test_workflow.py index f6cd1dada29..cdb4f08acc2 100644 --- a/envs/monkey_zoo/blackbox/tests/performance/performance_test_workflow.py +++ b/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() diff --git a/envs/monkey_zoo/blackbox/tests/performance/report_generation_from_telemetries.py b/envs/monkey_zoo/blackbox/tests/performance/report_generation_from_telemetries.py new file mode 100644 index 00000000000..a08bbda701c --- /dev/null +++ b/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() diff --git a/envs/monkey_zoo/blackbox/tests/performance/telemetry_performance_test_workflow.py b/envs/monkey_zoo/blackbox/tests/performance/telemetry_performance_test_workflow.py new file mode 100644 index 00000000000..b5acf4a9ea9 --- /dev/null +++ b/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() diff --git a/monkey/monkey_island/cc/app.py b/monkey/monkey_island/cc/app.py index 3a11349301b..13aac018aeb 100644 --- a/monkey/monkey_island/cc/app.py +++ b/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/', '/api/report//') + api.add_resource(ZeroTrustFindingEvent, '/api/zero-trust/finding-event/') api.add_resource(TelemetryFeed, '/api/telemetry-feed', '/api/telemetry-feed/') api.add_resource(Log, '/api/log', '/api/log/') diff --git a/monkey/monkey_island/cc/models/zero_trust/aggregate_finding.py b/monkey/monkey_island/cc/models/zero_trust/aggregate_finding.py index 38b5510478c..c3817313f5a 100644 --- a/monkey/monkey_island/cc/models/zero_trust/aggregate_finding.py +++ b/monkey/monkey_island/cc/models/zero_trust/aggregate_finding.py @@ -12,7 +12,7 @@ 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: @@ -21,7 +21,6 @@ def create_or_add_to_existing(test, status, events): # Now we know for sure this is the only one orig_finding = existing_findings[0] orig_finding.add_events(events) - orig_finding.save() def add_malicious_activity_to_timeline(events): diff --git a/monkey/monkey_island/cc/models/zero_trust/event.py b/monkey/monkey_island/cc/models/zero_trust/event.py index 89b581fa0e3..5ba909f2850 100644 --- a/monkey/monkey_island/cc/models/zero_trust/event.py +++ b/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, diff --git a/monkey/monkey_island/cc/models/zero_trust/finding.py b/monkey/monkey_island/cc/models/zero_trust/finding.py index 90c9e1dc3cc..2f3261ec4ac 100644 --- a/monkey/monkey_island/cc/models/zero_trust/finding.py +++ b/monkey/monkey_island/cc/models/zero_trust/finding.py @@ -2,6 +2,7 @@ """ Define a Document Schema for Zero Trust findings. """ +from typing import List from mongoengine import Document, StringField, EmbeddedDocumentListField @@ -55,6 +56,5 @@ def save_finding(test, status, events): return finding - def add_events(self, events): - # type: (list) -> None - self.events.extend(events) + def add_events(self, events: List) -> None: + self.update(push_all__events=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 index ce28ad7f7b3..b817f6e922f 100644 --- a/monkey/monkey_island/cc/models/zero_trust/test_aggregate_finding.py +++ b/monkey/monkey_island/cc/models/zero_trust/test_aggregate_finding.py @@ -1,3 +1,8 @@ +import unittest +from packaging import version + +import mongomock + import common.data.zero_trust_consts as zero_trust_consts from monkey_island.cc.models.zero_trust.aggregate_finding import AggregateFinding from monkey_island.cc.models.zero_trust.event import Event @@ -6,6 +11,9 @@ class TestAggregateFinding(IslandTestCase): + + @unittest.skipIf(version.parse(mongomock.__version__) <= version.parse("3.19.0"), + "mongomock version doesn't support this test") def test_create_or_add_to_existing(self): self.fail_if_not_testing_env() self.clean_finding_db() @@ -25,6 +33,8 @@ def test_create_or_add_to_existing(self): self.assertEqual(len(Finding.objects(test=test, status=status)), 1) self.assertEqual(len(Finding.objects(test=test, status=status)[0].events), 2) + @unittest.skipIf(version.parse(mongomock.__version__) <= version.parse("3.19.0"), + "mongomock version doesn't support this test") def test_create_or_add_to_existing_2_tests_already_exist(self): self.fail_if_not_testing_env() self.clean_finding_db() diff --git a/monkey/monkey_island/cc/resources/zero_trust/finding_event.py b/monkey/monkey_island/cc/resources/zero_trust/finding_event.py new file mode 100644 index 00000000000..16c5452410a --- /dev/null +++ b/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)} diff --git a/monkey/monkey_island/cc/services/reporting/test_zero_trust_service.py b/monkey/monkey_island/cc/services/reporting/test_zero_trust_service.py index 328be2e007d..e40af29f4e5 100644 --- a/monkey/monkey_island/cc/services/reporting/test_zero_trust_service.py +++ b/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._get_events_without_overlap(5, [1, 2, 3])) + self.assertListEqual([3], ZeroTrustService._get_events_without_overlap(6, [1, 2, 3])) + 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): t = list(t) # make a mutable copy diff --git a/monkey/monkey_island/cc/services/reporting/zero_trust_service.py b/monkey/monkey_island/cc/services/reporting/zero_trust_service.py index dd6fad1bc97..ee8fdd8bb49 100644 --- a/monkey/monkey_island/cc/services/reporting/zero_trust_service.py +++ b/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)) diff --git a/monkey/monkey_island/cc/testing/README.md b/monkey/monkey_island/cc/testing/README.md new file mode 100644 index 00000000000..1c1446b2f4f --- /dev/null +++ b/monkey/monkey_island/cc/testing/README.md @@ -0,0 +1,9 @@ +# Profiling island + +To profile specific methods on island a `@profile(sort_args=['cumulative'], print_args=[100])` +decorator can be used. +Use it as any other decorator. After decorated method is used, a file will appear in a +directory provided in `profiler_decorator.py`. Filename describes the path of +the method that was profiled. For example if method `monkey_island/cc/resources/netmap.get` +was profiled, then the results of this profiling will appear in +`monkey_island_cc_resources_netmap_get`. diff --git a/monkey/monkey_island/cc/testing/profiler_decorator.py b/monkey/monkey_island/cc/testing/profiler_decorator.py new file mode 100644 index 00000000000..997ef91aebc --- /dev/null +++ b/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(".", "_") diff --git a/monkey/monkey_island/cc/ui/src/components/report-components/zerotrust/EventsButton.js b/monkey/monkey_island/cc/ui/src/components/report-components/zerotrust/EventsButton.js index 62fc8e8e436..b7997057891 100644 --- a/monkey/monkey_island/cc/ui/src/components/report-components/zerotrust/EventsButton.js +++ b/monkey/monkey_island/cc/ui/src/components/report-components/zerotrust/EventsButton.js @@ -24,7 +24,12 @@ export default class EventsButton extends Component { render() { return -