feat: updated AssessmentKey

someengineering · Oct 31, 2024 · 09a097a · 09a097a
1 parent dc06447
commit 09a097a
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 5 deletions.
diff --git 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]]
 
@@ -509,7 +515,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
@@ -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