Improved threading for EC2

prowler/providers/aws/lib/service/service.py

@@ -1,5 +1,5 @@
 import threading
-
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from prowler.providers.aws.aws_provider import (
     generate_regional_clients,
     get_default_region,
@@ -47,11 +47,19 @@ class AWSService:
     def __get_session__(self):
         return self.session
 
-    def __threading_call__(self, call):
-        threads = []
-        for regional_client in self.regional_clients.values():
-            threads.append(threading.Thread(target=call, args=(regional_client,)))
-        for t in threads:
-            t.start()
-        for t in threads:
-            t.join()
+    def __threading_call__(self, call, iterator=None, max_workers=10):
+        # Use the provided iterator, or default to self.regional_clients
+        items = iterator if iterator is not None else self.regional_clients.values()
+
+        # Using ThreadPoolExecutor for managing threads
+        with ThreadPoolExecutor(max_workers=max_workers) as executor:
+            # Submit tasks to the executor
+            futures = [executor.submit(call, item) for item in items]
+
+            # Wait for all tasks to complete
+            for future in as_completed(futures):
+                try:
+                    future.result()  # Raises exceptions from the thread, if any
+                except Exception as e:
+                    # Handle exceptions if necessary
+                    pass  # Replace 'pass' with any additional exception handling logic

prowler/providers/aws/services/ec2/ec2_service.py

@@ -17,7 +17,7 @@ class EC2(AWSService):
         super().__init__(__class__.__name__, audit_info)
         self.instances = []
         self.__threading_call__(self.__describe_instances__)
-        self.__get_instance_user_data__()
+        self.__threading_call__(self.__get_instance_user_data__, self.instances)
         self.security_groups = []
         self.regions_with_sgs = []
         self.__threading_call__(self.__describe_security_groups__)
@@ -27,7 +27,7 @@ class EC2(AWSService):
         self.volumes_with_snapshots = {}
         self.regions_with_snapshots = {}
         self.__threading_call__(self.__describe_snapshots__)
-        self.__get_snapshot_public__()
+        self.__threading_call__(self.__get_snapshot_public__, self.snapshots)
         self.network_interfaces = []
         self.__threading_call__(self.__describe_public_network_interfaces__)
         self.__threading_call__(self.__describe_sg_network_interfaces__)
@@ -219,32 +219,29 @@ class EC2(AWSService):
                 f"{regional_client.region} -- {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
             )
 
-    def __get_snapshot_public__(self):
+    def __get_snapshot_public__(self, snapshot):
         logger.info("EC2 - Getting snapshot volume attribute permissions...")
-        for snapshot in self.snapshots:
-            try:
-                regional_client = self.regional_clients[snapshot.region]
-                snapshot_public = regional_client.describe_snapshot_attribute(
-                    Attribute="createVolumePermission", SnapshotId=snapshot.id
-                )
-                for permission in snapshot_public["CreateVolumePermissions"]:
-                    if "Group" in permission:
-                        if permission["Group"] == "all":
-                            snapshot.public = True
+        try:
+            regional_client = self.regional_clients[snapshot.region]
+            snapshot_public = regional_client.describe_snapshot_attribute(
+                Attribute="createVolumePermission", SnapshotId=snapshot.id
+            )
+            for permission in snapshot_public["CreateVolumePermissions"]:
+                if "Group" in permission:
+                    if permission["Group"] == "all":
+                        snapshot.public = True
 
-            except ClientError as error:
-                if error.response["Error"]["Code"] == "InvalidSnapshot.NotFound":
-                    logger.warning(
-                        f"{snapshot.region} --"
-                        f" {error.__class__.__name__}[{error.__traceback__.tb_lineno}]:"
-                        f" {error}"
-                    )
-                    continue
-
-            except Exception as error:
-                logger.error(
-                    f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
+        except ClientError as error:
+            if error.response["Error"]["Code"] == "InvalidSnapshot.NotFound":
+                logger.warning(
+                    f"{snapshot.region} --"
+                    f" {error.__class__.__name__}[{error.__traceback__.tb_lineno}]:"
+                    f" {error}"
                 )
+        except Exception as error:
+            logger.error(
+                f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
+            )
 
     def __describe_public_network_interfaces__(self, regional_client):
         logger.info("EC2 - Describing Network Interfaces...")
@@ -299,27 +296,24 @@ class EC2(AWSService):
                 f"{regional_client.region} -- {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
             )
 
-    def __get_instance_user_data__(self):
+    def __get_instance_user_data__(self, instance):
         logger.info("EC2 - Getting instance user data...")
-        for instance in self.instances:
-            try:
-                regional_client = self.regional_clients[instance.region]
-                user_data = regional_client.describe_instance_attribute(
-                    Attribute="userData", InstanceId=instance.id
-                )["UserData"]
-                if "Value" in user_data:
-                    instance.user_data = user_data["Value"]
-
-            except ClientError as error:
-                if error.response["Error"]["Code"] == "InvalidInstanceID.NotFound":
-                    logger.warning(
-                        f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
-                    )
-                    continue
-            except Exception as error:
-                logger.error(
+        try:
+            regional_client = self.regional_clients[instance.region]
+            user_data = regional_client.describe_instance_attribute(
+                Attribute="userData", InstanceId=instance.id
+            )["UserData"]
+            if "Value" in user_data:
+                instance.user_data = user_data["Value"]
+        except ClientError as error:
+            if error.response["Error"]["Code"] == "InvalidInstanceID.NotFound":
+                logger.warning(
                     f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
                 )
+        except Exception as error:
+            logger.error(
+                f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
+            )
 
     def __describe_images__(self, regional_client):
         logger.info("EC2 - Describing Images...")