diff --git a/monkey/common/utils/exceptions.py b/monkey/common/utils/exceptions.py index b1811ab88..333db5d60 100644 --- a/monkey/common/utils/exceptions.py +++ b/monkey/common/utils/exceptions.py @@ -36,3 +36,7 @@ class NoInternetError(Exception): class ScoutSuiteScanError(Exception): """ Raise to indicate problems ScoutSuite encountered during scanning""" + + +class UnknownFindingError(Exception): + """ Raise when provided finding is of unknown type""" diff --git a/monkey/monkey_island/cc/services/zero_trust/finding_service.py b/monkey/monkey_island/cc/services/zero_trust/finding_service.py index 618be9550..02459eb0e 100644 --- a/monkey/monkey_island/cc/services/zero_trust/finding_service.py +++ b/monkey/monkey_island/cc/services/zero_trust/finding_service.py @@ -1,6 +1,7 @@ from typing import List from common.common_consts import zero_trust_consts +from common.utils.exceptions import UnknownFindingError from monkey_island.cc.models.zero_trust.finding import Finding from monkey_island.cc.services.zero_trust.monkey_details_service import MonkeyDetailsService @@ -10,19 +11,20 @@ class FindingService: @staticmethod def get_all_findings() -> List[Finding]: findings = list(Finding.objects) - details = [] for i in range(len(findings)): if findings[i].finding_type == zero_trust_consts.MONKEY_FINDING: details = MonkeyDetailsService.fetch_details_for_display(findings[i].details.id) elif findings[i].finding_type == zero_trust_consts.SCOUTSUITE_FINDING: details = findings[i].details.fetch().to_mongo() + else: + raise UnknownFindingError(f"Unknown finding type {findings[i].finding_type}") findings[i] = findings[i].to_mongo() - findings[i] = FindingService.get_enriched_finding(findings[i]) + findings[i] = FindingService._get_enriched_finding(findings[i]) findings[i]['details'] = details return findings @staticmethod - def get_enriched_finding(finding): + def _get_enriched_finding(finding): test_info = zero_trust_consts.TESTS_MAP[finding['test']] enriched_finding = { 'finding_id': str(finding['_id']), diff --git a/monkey/monkey_island/cc/services/zero_trust/monkey_details_service.py b/monkey/monkey_island/cc/services/zero_trust/monkey_details_service.py index 623285ddd..5332ed90b 100644 --- a/monkey/monkey_island/cc/services/zero_trust/monkey_details_service.py +++ b/monkey/monkey_island/cc/services/zero_trust/monkey_details_service.py @@ -12,7 +12,7 @@ EVENT_FETCH_CNT = 50 class MonkeyDetailsService: @staticmethod - def fetch_details_for_display(finding_id: ObjectId): + def fetch_details_for_display(finding_id: ObjectId) -> dict: pipeline = [{'$match': {'_id': finding_id}}, {'$addFields': {'oldest_events': {'$slice': ['$events', EVENT_FETCH_CNT]}, 'latest_events': {'$slice': ['$events', -1 * EVENT_FETCH_CNT]},