diff --git a/backfill.py b/backfill.py index bd18223..e37d021 100644 --- a/backfill.py +++ b/backfill.py @@ -19,16 +19,32 @@ def get_log_groups(token=None): if token is None: - return cloudwatch_logs_client.describe_log_groups( - logGroupNamePrefix=log_group_prefix, limit=log_groups_return_limit - ) + if log_group_prefix != "": + return cloudwatch_logs_client.describe_log_groups( + logGroupNamePrefix=log_group_prefix, limit=log_groups_return_limit + ) + else: + return cloudwatch_logs_client.describe_log_groups( + limit=log_groups_return_limit + ) else: - return cloudwatch_logs_client.describe_log_groups( - logGroupNamePrefix=log_group_prefix, - nextToken=token, - limit=log_groups_return_limit, - ) + if log_group_prefix != "": + return cloudwatch_logs_client.describe_log_groups( + logGroupNamePrefix=log_group_prefix, + nextToken=token, + limit=log_groups_return_limit, + ) + else: + return cloudwatch_logs_client.describe_log_groups( + nextToken=token, + limit=log_groups_return_limit, + ) +def remove_permission(lambda_arn): + lambda_client.remove_permission( + FunctionName=lambda_arn, + StatementId="cloudwatch-backfiller-axiom", + ) def delete_subscription_filter(log_group_arn, lambda_arn): try: @@ -36,11 +52,6 @@ def delete_subscription_filter(log_group_arn, lambda_arn): logger.info(f"Deleting subscription filter for {log_group_name}...") - lambda_client.remove_permission( - FunctionName=lambda_arn, - StatementId="%s-axiom" % log_group_name.replace("/", "-"), - ) - cloudwatch_logs_client.delete_subscription_filter( logGroupName=log_group_name, filterName="%s-axiom" % log_group_name ) @@ -54,18 +65,25 @@ def delete_subscription_filter(log_group_arn, lambda_arn): raise e -def create_subscription_filter(log_group_arn, lambda_arn): - try: - log_group_name = log_group_arn.split(":")[-2] - logger.info(f"Creating subscription filter for {log_group_name}...") - lambda_client.add_permission( +def create_statement(region, account_id, lambda_arn): + logger.info(f"Creating permission for {lambda_arn}...") + source_arn = "arn:aws:logs:%s:%s:log-group:*:*" % ( + region, + account_id, + ) + lambda_client.add_permission( FunctionName=lambda_arn, - StatementId="%s-axiom" % log_group_name.replace("/", "-"), + StatementId="cloudwatch-backfiller-axiom", Action="lambda:InvokeFunction", Principal=f"logs.amazonaws.com", - SourceArn=log_group_arn, + SourceArn=source_arn, ) +def create_subscription_filter(log_group_arn, lambda_arn): + try: + log_group_name = log_group_arn.split(":")[-2] + logger.info(f"Creating subscription filter for {log_group_name}...") + cloudwatch_logs_client.put_subscription_filter( logGroupName=log_group_name, filterName="%s-axiom" % log_group_name, @@ -84,6 +102,19 @@ def create_subscription_filter(log_group_arn, lambda_arn): def lambda_handler(event: dict, context=None): if axiom_cloudwatch_lambda_ingester_arn is None: raise Exception("AXIOM_CLOUDWATCH_LAMBDA_INGESTER_ARN is not set") + + aws_account_id = context.invoked_function_arn.split(":")[4] + region = os.getenv("AWS_REGION") + + # create permission for lambda + try: + remove_permission(axiom_cloudwatch_lambda_ingester_arn) + except Exception as e: + logger.error(f"Error removing permission: {e}") + + create_statement(region, aws_account_id, axiom_cloudwatch_lambda_ingester_arn) + + ingester_lambda_group_name = "/aws/lambda/" + axiom_cloudwatch_lambda_ingester_arn.split(":")[-1] def log_groups(token=None): groups_response = get_log_groups(token) @@ -92,18 +123,27 @@ def log_groups(token=None): if len(groups) == 0: return - + for group in groups: + # skip the ingester lambda log group to avoid circular logging + if group["logGroupName"] == ingester_lambda_group_name: + continue + try: delete_subscription_filter( - group["arn"], axiom_cloudwatch_lambda_ingester_arn + region, aws_account_id, axiom_cloudwatch_lambda_ingester_arn ) except Exception: pass - create_subscription_filter( - group["arn"], axiom_cloudwatch_lambda_ingester_arn - ) + try: + create_subscription_filter( + group["arn"], axiom_cloudwatch_lambda_ingester_arn + ) + except cloudwatch_logs_client.exceptions.LimitExceededException as error: + logger.error("failed to create subscription filter for: %s" % group["logGroupName"]) + logger.error(error) + continue if token is None: return @@ -118,7 +158,13 @@ def log_groups(token=None): log_groups() except Exception as e: responseData["success"] = "False" - cfnresponse.send(event, context, cfnresponse.FAILED, responseData) + if "ResponseURL" in event: + cfnresponse.send(event, context, cfnresponse.FAILED, responseData) + else: + raise e responseData["success"] = "True" - cfnresponse.send(event, context, cfnresponse.SUCCESS, responseData) + if "ResponseURL" in event: + cfnresponse.send(event, context, cfnresponse.SUCCESS, responseData) + else: + return 'ok'