diff --git a/plugins/aws/fix_plugin_aws/resource/base.py b/plugins/aws/fix_plugin_aws/resource/base.py index 93b1064cbf..7aa819ab8d 100644 --- a/plugins/aws/fix_plugin_aws/resource/base.py +++ b/plugins/aws/fix_plugin_aws/resource/base.py @@ -65,8 +65,14 @@ def get_client(config: Config, resource: BaseResource) -> AwsClient: "region_id": lambda n: n.region().id, } -# Type alias for a tuple representing provider, region, and resource type -AssessmentKey = Tuple[str, str, str] + +@define(slots=True, frozen=True) +class AssessmentKey: + provider: str + region: str + resource_type: str + + # Type alias for the inner dictionary that maps resource ID to a list of findings ResourceFindings = Dict[str, List[Finding]] @@ -511,7 +517,7 @@ def suppress(self, message: str) -> SuppressWithFeedback: return SuppressWithFeedback(message, self.core_feedback, log) def add_finding(self, provider: str, class_name: str, region: str, class_id: str, finding: Finding) -> None: - self._assessment_findings[(provider, region, class_name)][class_id].append(finding) + self._assessment_findings[AssessmentKey(provider, region, class_name)][class_id].append(finding) def submit_work(self, service: str, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: """ diff --git a/plugins/aws/fix_plugin_aws/resource/guardduty.py b/plugins/aws/fix_plugin_aws/resource/guardduty.py index 11ca119654..bd4162439f 100644 --- a/plugins/aws/fix_plugin_aws/resource/guardduty.py +++ b/plugins/aws/fix_plugin_aws/resource/guardduty.py @@ -6,7 +6,7 @@ from attrs import define, field from boto3.exceptions import Boto3Error -from fix_plugin_aws.resource.base import AwsResource, GraphBuilder +from fix_plugin_aws.resource.base import AssessmentKey, AwsResource, GraphBuilder from fixlib.baseresources import Assessment, Finding, PhantomBaseResource, Severity from fixlib.json_bender import F, S, AsInt, Bend, Bender, ForallBend @@ -1278,7 +1278,7 @@ def set_findings(builder: GraphBuilder, resource_to_set: AwsResource, to_check: else: return provider_findings = builder._assessment_findings.get( - ("guard_duty", resource_to_set.region().id, resource_to_set.__class__.__name__), {} + AssessmentKey("guard_duty", resource_to_set.region().id, resource_to_set.__class__.__name__), {} ).get(id_or_arn_or_name, []) if provider_findings: # Set the findings in the resource's _assessments dictionary