From 5dfd8460be32a069d027bca0a24002cfb54bc359 Mon Sep 17 00:00:00 2001 From: Fennerr Date: Mon, 11 Dec 2023 11:06:32 +0200 Subject: [PATCH] Improved threading for EC2 --- prowler/providers/aws/lib/service/service.py | 26 +++--- .../providers/aws/services/ec2/ec2_service.py | 80 +++++++++---------- 2 files changed, 54 insertions(+), 52 deletions(-) diff --git a/prowler/providers/aws/lib/service/service.py b/prowler/providers/aws/lib/service/service.py index 0badb1a88d..43cef6b99a 100644 --- a/prowler/providers/aws/lib/service/service.py +++ b/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 diff --git a/prowler/providers/aws/services/ec2/ec2_service.py b/prowler/providers/aws/services/ec2/ec2_service.py index 4838f04115..21a37ca2d1 100644 --- a/prowler/providers/aws/services/ec2/ec2_service.py +++ b/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...")