Island: Fix mypy issues in finding_service.py

This commit is contained in:
Kekoa Kaaikala 2022-09-21 18:23:33 +00:00
parent 7013963d59
commit a5b5449f73
1 changed files with 15 additions and 10 deletions

View File

@ -1,5 +1,5 @@
from dataclasses import dataclass from dataclasses import dataclass
from typing import List, Union from typing import Dict, List, Union, cast
from bson import SON from bson import SON
@ -30,21 +30,26 @@ class FindingService:
@staticmethod @staticmethod
def get_all_findings_for_ui() -> List[EnrichedFinding]: def get_all_findings_for_ui() -> List[EnrichedFinding]:
findings = FindingService.get_all_findings_from_db() findings = FindingService.get_all_findings_from_db()
for i in range(len(findings)): enriched_findings: List[EnrichedFinding] = []
details = FindingService._get_finding_details(findings[i]) for finding in findings:
findings[i] = findings[i].to_mongo() finding_data = finding.to_mongo()
findings[i] = FindingService._get_enriched_finding(findings[i]) enriched_finding = FindingService._get_enriched_finding(finding_data)
findings[i].details = details details = FindingService._get_finding_details(finding)
return findings enriched_finding.details = details
enriched_findings.append(enriched_finding)
return enriched_findings
@staticmethod @staticmethod
def _get_enriched_finding(finding: Finding) -> EnrichedFinding: def _get_enriched_finding(finding: SON) -> EnrichedFinding:
test_info = zero_trust_consts.TESTS_MAP[finding["test"]] test_info = zero_trust_consts.TESTS_MAP[finding["test"]]
enriched_finding = EnrichedFinding( enriched_finding = EnrichedFinding(
finding_id=str(finding["_id"]), finding_id=str(finding["_id"]),
test=test_info[zero_trust_consts.FINDING_EXPLANATION_BY_STATUS_KEY][finding["status"]], test=cast(
Dict[str, str], test_info[zero_trust_consts.FINDING_EXPLANATION_BY_STATUS_KEY]
)[finding["status"]],
test_key=finding["test"], test_key=finding["test"],
pillars=test_info[zero_trust_consts.PILLARS_KEY], pillars=cast(List[str], test_info[zero_trust_consts.PILLARS_KEY]),
status=finding["status"], status=finding["status"],
details=None, details=None,
) )