From 98031fc14ab59a7d5325b5567973441c12f56ceb Mon Sep 17 00:00:00 2001 From: Charles Dunda Date: Tue, 10 Oct 2023 14:35:41 -0400 Subject: [PATCH 1/2] Add ACM Cert Handler --- .../aws/resources/acm_cert_handler.py | 80 +++++++++++++++++++ .../aws/resources/handler_creator.py | 9 ++- 2 files changed, 87 insertions(+), 2 deletions(-) create mode 100644 lambda_function/aws/resources/acm_cert_handler.py diff --git a/lambda_function/aws/resources/acm_cert_handler.py b/lambda_function/aws/resources/acm_cert_handler.py new file mode 100644 index 0000000..188de90 --- /dev/null +++ b/lambda_function/aws/resources/acm_cert_handler.py @@ -0,0 +1,80 @@ +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 ACMHandler(BaseHandler): + def handle(self): + for region in list(self.regions): + aws_acm_client = boto3.client("acm", region_name=region) + try: + paginator = aws_acm_client.get_paginator('list_certificates') + for page in paginator.paginate(): + certificate_summary_list = page.get('CertificateSummaryList', []) + self._handle_list_response(certificate_summary_list, region) + + 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) + except Exception as e: + logger.error(f"Failed to list ACM certificates in region: {region}; {e}") + self.skip_delete = True + + return { + "aws_entities": self.aws_entities, + "next_resource_config": None, + "skip_delete": self.skip_delete, + } + + def _handle_list_response(self, certificate_summary_list, region): + with ThreadPoolExecutor(max_workers=consts.MAX_CC_WORKERS) as executor: + futures = [ + executor.submit(self.handle_single_resource_item, region, cert.get("CertificateArn", "")) + for cert in certificate_summary_list + ] + 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_close_to_timeout(self, region): + 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, + } + + def handle_single_resource_item(self, region, certificate_arn, action_type="upsert"): + entities = [] + skip_delete = False + try: + resource_obj = {} + if action_type == "upsert": + logger.info(f"Get ACM certificate details for ARN: {certificate_arn}") + aws_acm_client = boto3.client("acm", region_name=region) + response = aws_acm_client.describe_certificate(CertificateArn=certificate_arn) + resource_obj = response.get("Certificate", {}) + + # Handles unserializable date properties in the JSON by turning them into a string + resource_obj = json.loads(json.dumps(resource_obj, default=str)) + elif action_type == "delete": + resource_obj = {"identifier": certificate_arn} # Entity identifier to delete + entities = create_entities_json(resource_obj, self.selector_query, self.mappings, action_type) + except Exception as e: + logger.error(f"Failed to extract or transform certificate ARN: {certificate_arn}, error: {e}") + skip_delete = True + + aws_entities = handle_entities(entities, self.port_client, action_type) + + return {"aws_entities": aws_entities, "skip_delete": skip_delete} diff --git a/lambda_function/aws/resources/handler_creator.py b/lambda_function/aws/resources/handler_creator.py index 558d3a8..413b6bf 100644 --- a/lambda_function/aws/resources/handler_creator.py +++ b/lambda_function/aws/resources/handler_creator.py @@ -5,9 +5,14 @@ from aws.resources.cloudformation_handler import CloudFormationHandler from aws.resources.ec2_instance_handler import EC2InstanceHandler from aws.resources.load_balancer_handler import LoadBalancerHandler +from aws.resources.acm_cert_handler import ACMHandler -SPECIAL_AWS_HANDLERS: Dict[str, Type[BaseHandler]] = {"AWS::CloudFormation::Stack": CloudFormationHandler, "AWS::EC2::Instance": EC2InstanceHandler, - "AWS::ElasticLoadBalancingV2::LoadBalancer": LoadBalancerHandler} +SPECIAL_AWS_HANDLERS: Dict[str, Type[BaseHandler]] = { + "AWS::CloudFormation::Stack": CloudFormationHandler, + "AWS::EC2::Instance": EC2InstanceHandler, + "AWS::ElasticLoadBalancingV2::LoadBalancer": LoadBalancerHandler, + "AWS::ACM::Certificate": ACMHandler, +} def create_resource_handler(resource_config, port_client, lambda_context, default_region): From e49b9e72e963cc4ef8d6a55c1e314ef5ec652795 Mon Sep 17 00:00:00 2001 From: Charles Dunda Date: Wed, 25 Oct 2023 16:41:00 -0400 Subject: [PATCH 2/2] fixup: update acm handler with pr review requests - Use the NextToken Pattern - allow list_parameters to allow other filters - add log for list with current region - add next_token self._cleanup_regions --- .../aws/resources/acm_cert_handler.py | 79 +++++++++++-------- 1 file changed, 44 insertions(+), 35 deletions(-) diff --git a/lambda_function/aws/resources/acm_cert_handler.py b/lambda_function/aws/resources/acm_cert_handler.py index 188de90..ef773af 100644 --- a/lambda_function/aws/resources/acm_cert_handler.py +++ b/lambda_function/aws/resources/acm_cert_handler.py @@ -9,52 +9,46 @@ logger = logging.getLogger(__name__) - class ACMHandler(BaseHandler): def handle(self): for region in list(self.regions): aws_acm_client = boto3.client("acm", region_name=region) - try: - paginator = aws_acm_client.get_paginator('list_certificates') - for page in paginator.paginate(): - certificate_summary_list = page.get('CertificateSummaryList', []) - self._handle_list_response(certificate_summary_list, region) + logger.info(f"List ACM Certificates, region: {region}") + self.next_token = "" if self.next_token is None else self.next_token + while self.next_token is not None: + list_certificates_params = self.selector_aws.get("list_parameters", {}) + if self.next_token: + list_certificates_params["NextToken"] = self.next_token + try: + response = aws_acm_client.list_certificates(**list_certificates_params) + except Exception as e: + logger.error( + f"Failed to list ACM Certificates, region: {region}," + f" Parameters: {list_certificates_params}; {e}") + self.skip_delete = True + self.next_token = None + break + + self._handle_list_response(response, region) + + self.next_token = response.get("NextToken") + 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) - 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) - except Exception as e: - logger.error(f"Failed to list ACM certificates in region: {region}; {e}") - self.skip_delete = True + self._cleanup_regions(region) - return { - "aws_entities": self.aws_entities, - "next_resource_config": None, - "skip_delete": self.skip_delete, - } + return {"aws_entities": self.aws_entities, "next_resource_config": None, "skip_delete": self.skip_delete} - def _handle_list_response(self, certificate_summary_list, region): - with ThreadPoolExecutor(max_workers=consts.MAX_CC_WORKERS) as executor: - futures = [ - executor.submit(self.handle_single_resource_item, region, cert.get("CertificateArn", "")) - for cert in certificate_summary_list - ] + def _handle_list_response(self, list_response, region): + certificates = list_response.get("CertificateSummaryList", []) + with ThreadPoolExecutor(max_workers=consts.MAX_DEFAULT_AWS_WORKERS) as executor: + futures = [executor.submit(self.handle_single_resource_item, region, cert.get("CertificateArn")) for cert in certificates] 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_close_to_timeout(self, region): - 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, - } - def handle_single_resource_item(self, region, certificate_arn, action_type="upsert"): entities = [] skip_delete = False @@ -65,7 +59,7 @@ def handle_single_resource_item(self, region, certificate_arn, action_type="upse aws_acm_client = boto3.client("acm", region_name=region) response = aws_acm_client.describe_certificate(CertificateArn=certificate_arn) resource_obj = response.get("Certificate", {}) - + # Handles unserializable date properties in the JSON by turning them into a string resource_obj = json.loads(json.dumps(resource_obj, default=str)) elif action_type == "delete": @@ -78,3 +72,18 @@ def handle_single_resource_item(self, region, certificate_arn, action_type="upse 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}