fix-recursion-associations (#58)

lib/context/context.py

@@ -85,6 +85,7 @@ def get_context_config(self):
         # If there are no filters, we forced to return True as we expected a Match always
         resource_matched = False if self.mh_filters_config else True
         resource_config = False
+        all_associations = {}
 
         # If the resources lives in another account, we need the --mh-assume-role
         if self.resource_account_id != self.current_account_id and not self.sess:
@@ -94,17 +95,19 @@ def get_context_config(self):
                 self.resource_account_id,
                 self.current_account_id,
             )
-            return resource_config, resource_matched
+            return resource_config, resource_matched, all_associations
 
         # Get Handler
         hnld = self.get_handler()
         if not hnld:
-            return resource_config, resource_matched
+            return resource_config, resource_matched, all_associations
 
         # Execute Drilled
         if self.drilled_down:
             try:
-                hnld.execute_drilled_metachecks(self.cached_associated_resources)
+                all_associations = hnld.execute_drilled_metachecks(
+                    self.cached_associated_resources
+                )
             except (AttributeError, Exception) as err:
                 if "should return None" in str(err):
                     self.logger.info(
@@ -138,7 +141,7 @@ def get_context_config(self):
                 err,
             )
 
-        return resource_config, resource_matched
+        return resource_config, resource_matched, all_associations
 
     def get_context_tags(self):
         self.logger.info(

lib/context/resources/Base.py

@@ -124,6 +124,7 @@ def check_associated_resources(resource, level):
                         r, AwsIamUserMetacheck
                     )
                     resource.iam_users[r] = resource_drilled_output
+                    self.all_associations[r] = resource_drilled_output
                     if level < 1 and resource_drilled:
                         check_associated_resources(resource_drilled, level + 1)
 
@@ -142,6 +143,7 @@ def check_associated_resources(resource, level):
                         r, SecurityGroupMetacheck
                     )
                     resource.security_groups[r] = resource_drilled_output
+                    self.all_associations[r] = resource_drilled_output
                     if level < 1 and resource_drilled:
                         check_associated_resources(resource_drilled, level + 1)
 
@@ -160,7 +162,8 @@ def check_associated_resources(resource, level):
                         r, AwsIamRoleMetaCheck
                     )
                     resource.iam_roles[r] = resource_drilled_output
-                    if level < 2 and resource_drilled:
+                    self.all_associations[r] = resource_drilled_output
+                    if level < 1 and resource_drilled:
                         check_associated_resources(resource_drilled, level + 1)
 
             # IAM Policies
@@ -178,6 +181,7 @@ def check_associated_resources(resource, level):
                         r, IamPolicyMetacheck
                     )
                     resource.iam_policies[r] = resource_drilled_output
+                    self.all_associations[r] = resource_drilled_output
                     if level < 1 and resource_drilled:
                         check_associated_resources(resource_drilled, level + 1)
 
@@ -196,6 +200,7 @@ def check_associated_resources(resource, level):
                         r, AwsAutoScalingAutoScalingGroupMetacheck
                     )
                     resource.autoscaling_groups[r] = resource_drilled_output
+                    self.all_associations[r] = resource_drilled_output
                     if level < 1 and resource_drilled:
                         check_associated_resources(resource_drilled, level + 1)
 
@@ -214,6 +219,7 @@ def check_associated_resources(resource, level):
                         r, VolumeMetacheck
                     )
                     resource.volumes[r] = resource_drilled_output
+                    self.all_associations[r] = resource_drilled_output
                     if level < 1 and resource_drilled:
                         check_associated_resources(resource_drilled, level + 1)
 
@@ -228,6 +234,7 @@ def check_associated_resources(resource, level):
                 for r, v in list(resource.vpcs.items()):
                     resource_drilled_output, resource_drilled = execute(r, VpcMetacheck)
                     resource.vpcs[r] = resource_drilled_output
+                    self.all_associations[r] = resource_drilled_output
                     if level < 1 and resource_drilled:
                         check_associated_resources(resource_drilled, level + 1)
 
@@ -246,6 +253,7 @@ def check_associated_resources(resource, level):
                         r, SubnetMetacheck
                     )
                     resource.subnets[r] = resource_drilled_output
+                    self.all_associations[r] = resource_drilled_output
                     if level < 1 and resource_drilled:
                         check_associated_resources(resource_drilled, level + 1)
 
@@ -264,6 +272,7 @@ def check_associated_resources(resource, level):
                         r, RouteTableMetacheck
                     )
                     resource.route_tables[r] = resource_drilled_output
+                    self.all_associations[r] = resource_drilled_output
                     if level < 1 and resource_drilled:
                         check_associated_resources(resource_drilled, level + 1)
 
@@ -282,6 +291,7 @@ def check_associated_resources(resource, level):
                         r, ApiGatewayV2ApiMetacheck
                     )
                     resource.api_gwv2_apis[r] = resource_drilled_output
+                    self.all_associations[r] = resource_drilled_output
                     if level < 1 and resource_drilled:
                         check_associated_resources(resource_drilled, level + 1)
 
@@ -300,11 +310,15 @@ def check_associated_resources(resource, level):
                         r, AwsEc2InstanceMetacheck
                     )
                     resource.instances[r] = resource_drilled_output
+                    self.all_associations[r] = resource_drilled_output
                     if level < 1 and resource_drilled:
                         check_associated_resources(resource_drilled, level + 1)
 
+        self.all_associations = {}
         check_associated_resources(self, 0)
 
+        return self.all_associations
+
     def output_checks_drilled(self):
         mh_values_checks = {}
         context_config = self.checks()

lib/findings.py

@@ -242,10 +242,9 @@ def evaluate_finding(
             cached_associated_resources,
         )
         if "config" in context_options:
-            mh_config, mh_checks_matched = context.get_context_config()
-            # Get and Cache the associations for this resource
-            if mh_config:
-                cached_associated_resources.update(get_associations(mh_config))
+            mh_config, mh_checks_matched, all_association = context.get_context_config()
+            # Cache the associations for this resource
+            cached_associated_resources.update(all_association)
         else:
             mh_config = False
             mh_checks_matched = True
@@ -331,24 +330,3 @@ def evaluate_finding(
         mh_inventory,
         AwsAccountData,
     )
-
-
-# From each resource, get the associations, so we can cache them and avoid to get them again
-def get_associations(resource):
-    associations_all = {}
-
-    def get_associations_recursively(dictionary, parent_key=""):
-        for key, value in dictionary.items():
-            if isinstance(value, dict):
-                if key == "associations":
-                    for atype, associations in value.items():
-                        if isinstance(associations, dict):
-                            for association, association_values in associations.items():
-                                if association_values:
-                                    associations_all[association] = association_values
-                get_associations_recursively(
-                    value, f"{parent_key}.{key}" if parent_key else key
-                )
-
-    get_associations_recursively(resource)
-    return associations_all