update lambda-manager code to trigger on creation

src/lambda-manager/CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## lambda-manager
 
+## 2.0.2  / 24-06-2024
+### 💡 Enhancements 💡
+- Update the lambda to trigger on creation if ScanOldLogGroups is set to true
+
 ## 2.0.1  / 04-2-2024
 ### 💡 Enhancements 💡
 - Update lambda code so it will not require the allow all policy

src/lambda-manager/lambda_function.py

@@ -2,11 +2,13 @@
 import os
 import re
 import uuid
+import cfnresponse 
+
 def identify_arn_service(arn):
-    arn_parts = arn.split(':')    
+    arn_parts = arn.split(':')
     if len(arn_parts) < 6:
-        return "Invalid ARN format"    
-    service = arn_parts[2]    
+        return "Invalid ARN format"
+    service = arn_parts[2]
     if service == "lambda":
         return "lambda"
     elif service == "firehose":
@@ -25,10 +27,10 @@ def list_log_groups_and_subscriptions(cloudwatch_logs, regex_pattern, logs_filte
         response = cloudwatch_logs.describe_log_groups(**kwargs)
         log_groups.extend(response['logGroups'])
     for log_group in log_groups:
-        log_group_name = log_group['logGroupName'] 
+        log_group_name = log_group['logGroupName']
         if regex_pattern and re.match(regex_pattern, log_group_name):
             print(f"Log Group: {log_group_name}")
-            
+
             subscriptions = cloudwatch_logs.describe_subscription_filters(logGroupName=log_group_name)
             subscriptions = subscriptions.get('subscriptionFilters')
             print(f" Subscriptions: {subscriptions} for {log_group_name}")
@@ -121,75 +123,92 @@ def list_log_groups_and_subscriptions(cloudwatch_logs, regex_pattern, logs_filte
                 else:
                     print(f"Invalid destination type {destination_type}")
 
-
-
 def lambda_handler(event, context):
-    cloudwatch_logs = boto3.client('logs')
-    regex_pattern = os.environ.get('REGEX_PATTERN')
-    destination_type = os.environ.get('DESTINATION_TYPE')
-    logs_filter = os.environ.get('LOGS_FILTER', '')
-    scan_all_log_groups = os.environ.get('SCAN_OLD_LOGGROUPS', 'false')
-    destination_arn = os.environ.get('DESTINATION_ARN')
-    role_arn = os.environ.get('DESTINATION_ROLE')
-    filter_name = 'Coralogix_Filter_' + str(uuid.uuid4())
-    print(f"Scanning all log groups: {scan_all_log_groups}")
-    if scan_all_log_groups == 'true':
-        list_log_groups_and_subscriptions(cloudwatch_logs, regex_pattern, logs_filter, destination_arn, role_arn, filter_name, context)
-        lambda_client = boto3.client('lambda')
-        function_name = context.function_name
+    status = cfnresponse.SUCCESS
+    try:
+        cloudwatch_logs     = boto3.client('logs')
+        regex_pattern       = os.environ.get('REGEX_PATTERN')
+        destination_type    = os.environ.get('DESTINATION_TYPE')
+        logs_filter         = os.environ.get('LOGS_FILTER', '')
+        scan_all_log_groups = os.environ.get('SCAN_OLD_LOGGROUPS', 'false')
+        destination_arn     = os.environ.get('DESTINATION_ARN')
+        role_arn            = os.environ.get('DESTINATION_ROLE')
+        filter_name         = 'Coralogix_Filter_' + str(uuid.uuid4())
 
-        # Fetch the current function configuration
-        current_config = lambda_client.get_function_configuration(FunctionName=function_name)
-        current_env_vars = current_config['Environment']['Variables']
+        print(f"Scanning all log groups: {scan_all_log_groups}")
+        if scan_all_log_groups == 'true' and "RequestType" in event:
+            list_log_groups_and_subscriptions(cloudwatch_logs, regex_pattern, logs_filter, destination_arn, role_arn, filter_name, context)
+            lambda_client = boto3.client('lambda')
+            function_name = context.function_name
 
-        # Update the environment variables
-        current_env_vars['SCAN_OLD_LOGGROUPS'] = 'false'
+            # Fetch the current function configuration
+            current_config = lambda_client.get_function_configuration(FunctionName=function_name)
+            current_env_vars = current_config['Environment']['Variables']
 
-        # Update the Lambda function configuration
-        try:
-            response = lambda_client.update_function_configuration(
-                FunctionName=function_name,
-                Environment={'Variables': current_env_vars}
-            )
-            print("Updated environment variables:", response['Environment']['Variables'])
-        except Exception as e:
-            print("Error updating function configuration:", e)
-    else:
-        log_group_to_subscribe = event['detail']['requestParameters']['logGroupName']
-        print(f"Log Group: {log_group_to_subscribe}") 
-        if regex_pattern and re.match(regex_pattern, log_group_to_subscribe):           
-            if destination_type == 'firehose':
+            # Update the environment variables
+            current_env_vars['SCAN_OLD_LOGGROUPS'] = 'false'
+
+            # Update the Lambda function configuration
+            try:
+                response = lambda_client.update_function_configuration(
+                    FunctionName=function_name,
+                    Environment={'Variables': current_env_vars}
+                )
+                print("Updated environment variables:", response['Environment']['Variables'])
+            except Exception as e:
+                print("Error updating function configuration:", e)
+                status = cfnresponse.FAILED
+        elif scan_all_log_groups != 'true' and "RequestType" not in event:
+            log_group_to_subscribe = event['detail']['requestParameters']['logGroupName']
+            print(f"Log Group: {log_group_to_subscribe}")
+            if regex_pattern and re.match(regex_pattern, log_group_to_subscribe):
+                if destination_type == 'firehose':
+                        try:
+                            cloudwatch_logs.put_subscription_filter(
+                                destinationArn=destination_arn,
+                                roleArn=role_arn,
+                                filterName= filter_name,
+                                filterPattern=logs_filter,
+                                logGroupName=log_group_to_subscribe,
+                            )
+                        except Exception as e:
+                            print(f"Failed to put subscription filter for {log_group_to_subscribe}: {e}")
+                            status = cfnresponse.FAILED
+                elif destination_type == 'lambda':
                     try:
+                        lambda_client = boto3.client('lambda')
+                        region = context.invoked_function_arn.split(":")[3]
+                        account_id = context.invoked_function_arn.split(":")[4]
+                        lambda_client.add_permission(
+                          FunctionName=destination_arn,
+                          StatementId=f'allow-trigger-from-{log_group_to_subscribe}',
+                          Action='lambda:InvokeFunction',
+                          Principal='logs.amazonaws.com',
+                          SourceArn=f'arn:aws:logs:{region}:{account_id}:log-group:{log_group_to_subscribe}:*',
+                        )
                         cloudwatch_logs.put_subscription_filter(
                             destinationArn=destination_arn,
-                            roleArn=role_arn,
-                            filterName= filter_name,
+                            filterName= "coralogix-aws-shipper-cloudwatch-trigger",
                             filterPattern=logs_filter,
                             logGroupName=log_group_to_subscribe,
                         )
                     except Exception as e:
                         print(f"Failed to put subscription filter for {log_group_to_subscribe}: {e}")
-            elif destination_type == 'lambda':
-                try:
-                    lambda_client = boto3.client('lambda')
-                    region = context.invoked_function_arn.split(":")[3]
-                    account_id = context.invoked_function_arn.split(":")[4]
-                    lambda_client.add_permission(
-                      FunctionName=destination_arn,
-                      StatementId=f'allow-trigger-from-{log_group_to_subscribe}',
-                      Action='lambda:InvokeFunction',
-                      Principal='logs.amazonaws.com',
-                      SourceArn=f'arn:aws:logs:{region}:{account_id}:log-group:{log_group_to_subscribe}:*',
-                    )
-                    cloudwatch_logs.put_subscription_filter(
-                        destinationArn=destination_arn,
-                        filterName= "coralogix-aws-shipper-cloudwatch-trigger",
-                        filterPattern=logs_filter,
-                        logGroupName=log_group_to_subscribe,
-                    )
-                except Exception as e:
-                    print(f"Failed to put subscription filter for {log_group_to_subscribe}: {e}")
+                        status = cfnresponse.FAILED
+                else:
+                    print(f"Invalid destination type {destination_type}")
+                    status = cfnresponse.FAILED
             else:
-                print(f"Invalid destination type {destination_type}")  
-        else:
-            print(f"Loggroup {log_group_to_subscribe} excluded")
+                print(f"Loggroup {log_group_to_subscribe} excluded")
+    except Exception as e:
+        print(f"Failed with exception: {e}")
+        status = cfnresponse.FAILED
+    finally:
+        print("Sending response to custom resource")
+        cfnresponse.send(
+            event,
+            context,
+            status,
+            {},
+            event.get('PhysicalResourceId', context.aws_request_id)
+        )

src/lambda-manager/requirements.txt

@@ -0,0 +1 @@
+cfnresponse

src/lambda-manager/template.yaml

@@ -16,7 +16,7 @@ Metadata:
       - cloudwatch
       - lambda
     HomePageUrl: https://coralogix.com
-    SemanticVersion: 2.0.1
+    SemanticVersion: 2.0.2
     SourceCodeUrl: https://github.com/coralogix/coralogix-aws-serverless
   AWS::CloudFormation::Interface:
     ParameterGroups:
@@ -81,6 +81,10 @@ Conditions:
       - Fn::Equals:
           - Ref: NotificationEmail
           - ''
+  ScanOldLogGroups:
+    Fn::Equals:
+      - Ref: ScanOldLogGroups
+      - 'true'
 Resources:
   LambdaExecutionRole:
     Type: AWS::IAM::Role
@@ -175,6 +179,14 @@ Resources:
                   - "CreateLogGroup"
             Target:
                 Id: cx-loggroup-target
+
+  LambdaTrigger:
+    Type: Custom::LambdaTrigger
+    Condition: ScanOldLogGroups
+    DependsOn: LambdaFunction
+    Properties:
+      ServiceToken: !GetAtt LambdaFunction.Arn
+
   LambdaFunctionNotificationSubscription:
     Type: AWS::SNS::Subscription
     Condition: IsNotificationEnabled