diff --git a/monkey/monkey_island/cc/services/zero_trust/zero_trust_report/finding_service.py b/monkey/monkey_island/cc/services/zero_trust/zero_trust_report/finding_service.py index 2c5322d92..3f972bb4c 100644 --- a/monkey/monkey_island/cc/services/zero_trust/zero_trust_report/finding_service.py +++ b/monkey/monkey_island/cc/services/zero_trust/zero_trust_report/finding_service.py @@ -1,4 +1,7 @@ -from typing import List +from dataclasses import dataclass +from typing import List, Union + +from bson import SON from common.common_consts import zero_trust_consts from common.utils.exceptions import UnknownFindingError @@ -6,32 +9,48 @@ from monkey_island.cc.models.zero_trust.finding import Finding from monkey_island.cc.services.zero_trust.monkey_findings.monkey_zt_details_service import MonkeyZTDetailsService +@dataclass +class EnrichedFinding: + finding_id: str + test: str + test_key: str + pillars: List[str] + status: str + finding_type: str + details: Union[dict, None] + + class FindingService: @staticmethod - def get_all_findings() -> List[Finding]: + def get_all_findings() -> List[EnrichedFinding]: findings = list(Finding.objects) for i in range(len(findings)): - if findings[i].finding_type == zero_trust_consts.MONKEY_FINDING: - details = MonkeyZTDetailsService.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}") + details = FindingService._get_finding_details(findings[i]) findings[i] = findings[i].to_mongo() findings[i] = FindingService._get_enriched_finding(findings[i]) - findings[i]['details'] = details + findings[i].details = details return findings @staticmethod - def _get_enriched_finding(finding: Finding) -> dict: + def _get_enriched_finding(finding: Finding) -> EnrichedFinding: test_info = zero_trust_consts.TESTS_MAP[finding['test']] - enriched_finding = { - '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'], - 'finding_type': finding['finding_type'] - } + enriched_finding = EnrichedFinding( + 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'], + finding_type=finding['finding_type'], + details=None + ) return enriched_finding + + @staticmethod + def _get_finding_details(finding: Finding) -> Union[dict, SON]: + if finding.finding_type == zero_trust_consts.MONKEY_FINDING: + return MonkeyZTDetailsService.fetch_details_for_display(finding.details.id) + elif finding.finding_type == zero_trust_consts.SCOUTSUITE_FINDING: + return finding.details.fetch().to_mongo() + else: + raise UnknownFindingError(f"Unknown finding type {finding.finding_type}")