Support for multiple assume roles to get collective inventory

plugins/doc_fragments/assume_role.py

@@ -22,4 +22,10 @@ class ModuleDocFragment:
       - The ARN of the IAM role to assume to perform the lookup.
       - You should still provide AWS credentials with enough privilege to perform the AssumeRole action.
     aliases: ["iam_role_arn"]
+  assume_role_arns:
+    description:
+      - List of ARNs which can be assumed to get collective inventory result.
+      - This conflicts with `assume_role_arn`.
+      - You should still provide AWS credentials with enough privilege to perform the AssumeRole action.
+    aliases: ["iam_role_arns"]
 """

plugins/inventory/aws_ec2.py

@@ -611,7 +611,7 @@ def _get_all_hostnames(self, instance, hostnames):
 
         return hostname_list
 
-    def _query(self, regions, include_filters, exclude_filters, strict_permissions, use_ssm_inventory):
+    def _query_instances(self, regions, include_filters, exclude_filters, strict_permissions, use_ssm_inventory):
         """
         :param regions: a list of regions to query
         :param include_filters: a list of boto3 filter dictionaries
@@ -621,6 +621,7 @@ def _query(self, regions, include_filters, exclude_filters, strict_permissions,
         """
         instances = []
         ids_to_ignore = []
+
         for filter in exclude_filters:
             for i in self._get_instances_by_region(
                 regions,
@@ -644,7 +645,26 @@ def _query(self, regions, include_filters, exclude_filters, strict_permissions,
             for connection, _region in self.all_clients("ssm"):
                 self._add_ssm_information(connection, instances)
 
-        return {"aws_ec2": instances}
+        return instances
+
+    def _query(self, regions, include_filters, exclude_filters, strict_permissions, use_ssm_inventory, assume_role_arns):
+        """
+        :param regions: a list of regions to query
+        :param include_filters: a list of boto3 filter dictionaries
+        :param exclude_filters: a list of boto3 filter dictionaries
+        :param strict_permissions: a boolean determining whether to fail or ignore 403 error codes
+        :param assume_role_arns: list of arns to assume to acquire collective inventory
+
+        """
+
+        if assume_role_arns:
+            instances = []
+            for role in assume_role_arns:
+                super().collective_roles(role)
+                instances.extend(self._query_instances(regions, include_filters, exclude_filters, strict_permissions, use_ssm_inventory))
+            return {"aws_ec2": instances}
+        else:
+            return {"aws_ec2": self._query_instances(regions, include_filters, exclude_filters, strict_permissions, use_ssm_inventory)}
 
     def _add_ssm_information(self, connection, instances):
         filters = [{"Key": "AWS:InstanceInformation.InstanceId", "Values": [x["InstanceId"] for x in instances]}]
@@ -763,6 +783,7 @@ def parse(self, inventory, loader, path, cache=True):
             self._sanitize_group_name = self._legacy_script_compatible_group_sanitization
 
         # get user specifications
+        assume_role_arns = self.get_option("assume_role_arns")
         regions = self.get_option("regions")
         include_filters = self.build_include_filters()
         exclude_filters = self.get_option("exclude_filters")
@@ -788,7 +809,7 @@ def parse(self, inventory, loader, path, cache=True):
         result_was_cached, results = self.get_cached_result(path, cache)
 
         if not result_was_cached:
-            results = self._query(regions, include_filters, exclude_filters, strict_permissions, use_ssm_inventory)
+            results = self._query(regions, include_filters, exclude_filters, strict_permissions, use_ssm_inventory, assume_role_arns)
 
         self._populate(
             results,

plugins/plugin_utils/inventory.py

@@ -111,8 +111,13 @@ def _freeze_iam_role(self, iam_role_arn):
     def _set_frozen_credentials(self):
         options = self.get_options()
         iam_role_arn = options.get("assume_role_arn")
-        if iam_role_arn:
+        iam_role_arns = options.get("assume_role_arns")
+        if iam_role_arn and not iam_role_arns:
             self._freeze_iam_role(iam_role_arn)
+        if iam_role_arn and iam_role_arns:
+            self.fail_aws(
+                " Conflict here, cannot have not assume_role_arn and assume_role_arns used together "
+            )
 
     def _describe_regions(self, service):
         # Try pulling a list of regions from the service
@@ -160,6 +165,18 @@ def _boto3_regions(self, service):
 
         return regions
 
+    def collective_roles(self, role):
+        """
+        This is used when the user uses `assume_role_arns`, to get inventory from different accounts
+
+        :param role: The role mentioned in list of assume roles
+        """
+
+        # This will be resetting the credential and
+        # using the base credential to assume different roles
+        self._frozen_credentials = {}
+        self._freeze_iam_role(role)
+
     def all_clients(self, service):
         """
         Generator that yields a boto3 client and the region