Merge pull request #9 from PeyGis/elb

Export AWS Load Balancers to Port

lambda_function/aws/resources/handler_creator.py

@@ -4,8 +4,10 @@
 from aws.resources.cloudcontrol_handler import CloudControlHandler
 from aws.resources.cloudformation_handler import CloudFormationHandler
 from aws.resources.ec2_instance_handler import EC2InstanceHandler
+from aws.resources.load_balancer_handler import LoadBalancerHandler
 
-SPECIAL_AWS_HANDLERS: Dict[str, Type[BaseHandler]] = {"AWS::CloudFormation::Stack": CloudFormationHandler, "AWS::EC2::Instance": EC2InstanceHandler}
+SPECIAL_AWS_HANDLERS: Dict[str, Type[BaseHandler]] = {"AWS::CloudFormation::Stack": CloudFormationHandler, "AWS::EC2::Instance": EC2InstanceHandler,
+                                                      "AWS::ElasticLoadBalancingV2::LoadBalancer": LoadBalancerHandler}
 
 
 def create_resource_handler(resource_config, port_client, lambda_context, default_region):

lambda_function/aws/resources/load_balancer_handler.py

@@ -0,0 +1,108 @@
+import json
+import logging
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+import boto3
+import consts
+from aws.resources.base_handler import BaseHandler
+from port.entities import create_entities_json, handle_entities
+
+logger = logging.getLogger(__name__)
+
+
+class LoadBalancerHandler(BaseHandler):
+    def handle(self):
+        for region in list(self.regions):
+            aws_elbv2_client = boto3.client("elbv2", region_name=region)
+            logger.info(f"List Load Balancers, region: {region}")
+            self.next_token = "" if self.next_token is None else self.next_token
+            while self.next_token is not None:
+                filter_parameters = self.selector_aws.get("list_parameters", {})
+                if self.next_token:
+                    filter_parameters["Marker"] = self.next_token
+                try:
+                    response = aws_elbv2_client.describe_load_balancers(**filter_parameters)
+                except Exception as e:
+                    logger.error(
+                        f"Failed list Load Balancers, region: {region},"
+                        f" Parameters: {filter_parameters}; {e}")
+                    self.skip_delete = True
+                    self.next_token = None
+                    break
+
+                self._handle_list_response(response, region)
+
+                self.next_token = response.get("NextMarker")
+                if self.lambda_context.get_remaining_time_in_millis()< consts.REMAINING_TIME_TO_REINVOKE_THRESHOLD:
+                    # Lambda timeout is too close, should return checkpoint for next run
+                    return self._handle_close_to_timeout(region)
+
+            self._cleanup_regions(region)
+
+        return {"aws_entities": self.aws_entities, "next_resource_config": None, "skip_delete": self.skip_delete}
+
+    def _handle_list_response(self, list_response, region):
+        load_balancers = list_response.get("LoadBalancers", [])
+        with ThreadPoolExecutor(max_workers=consts.MAX_UPSERT_WORKERS) as executor:
+            futures = [executor.submit(self.handle_single_resource_item, region, load_balancer.get("LoadBalancerName")) for load_balancer in load_balancers]
+            for completed_future in as_completed(futures):
+                result = completed_future.result()
+                self.aws_entities.update(result.get("aws_entities", set()))
+                self.skip_delete = result.get("skip_delete", False) if not self.skip_delete else self.skip_delete
+
+    def handle_single_resource_item(self, region, elb_name, action_type="upsert"):
+        entities = []
+        skip_delete = False
+        try:
+            if action_type == "upsert":
+                logger.info(f"Get Load Balancer, Name: {elb_name} in {region}")
+
+                # Create a Boto3 client for the Elastic Load Balancing service
+                aws_elbv2_client = boto3.client('elbv2')
+                response = aws_elbv2_client.describe_load_balancers(Names=[elb_name])
+                # Extract load balancer details from the response
+                load_balancer_obj = response['LoadBalancers'][0]
+                load_balancer_arn = load_balancer_obj["LoadBalancerArn"]
+
+                ## Fetch the attributes attached to the load balancer
+                elb_attributes_response = aws_elbv2_client.describe_load_balancer_attributes(LoadBalancerArn=load_balancer_arn)
+                load_balancer_obj["Attributes"] = elb_attributes_response["Attributes"]
+
+                ## Fetch the listeners attatched to this load balancer
+                elb_listeners_response = aws_elbv2_client.describe_listeners(LoadBalancerArn=load_balancer_arn)
+                load_balancer_obj["Listeners"] = elb_listeners_response["Listeners"]
+
+                ## Fetch the tags attached to this load balancer
+                elb_tags_response = aws_elbv2_client.describe_tags(ResourceArns=[load_balancer_arn])
+                load_balancer_obj["Tags"] = elb_tags_response["TagDescriptions"][0]["Tags"]
+
+                # Handles unserializable date properties in the JSON by turning them into a string
+                load_balancer_obj = json.loads(json.dumps(load_balancer_obj, default=str))
+
+            elif action_type == "delete":
+                load_balancer_obj = {"identifier": elb_name}  # Entity identifier to delete
+
+            entities = create_entities_json(load_balancer_obj, self.selector_query, self.mappings, action_type)
+
+        except Exception as e:
+            logger.error(f"Failed to extract or transform Load Balancer with name: {elb_name}, error: {e}")
+            skip_delete = True
+
+        aws_entities = handle_entities(entities, self.port_client, action_type)
+
+        return {"aws_entities": aws_entities, "skip_delete": skip_delete}
+
+    def _handle_close_to_timeout(self, region):
+        if self.next_token:
+            self.selector_aws["next_token"] = self.next_token
+        else:
+            self.selector_aws.pop("next_token", None)
+            self._cleanup_regions(region)
+            if not self.regions:  # Nothing left to sync
+                return {"aws_entities": self.aws_entities, "next_resource_config": None, "skip_delete": self.skip_delete}
+
+        if "selector" not in self.resource_config:
+            self.resource_config["selector"] = {}
+        self.resource_config["selector"]["aws"] = self.selector_aws
+
+        return {"aws_entities": self.aws_entities, "next_resource_config": self.resource_config, "skip_delete": self.skip_delete}