import boto3 import json # Specify your SSO profile name here sso_profile = 'dbaas' def get_client(service, region): session = boto3.Session(profile_name=sso_profile, region_name=region) return session.client(service) def get_all_regions(service_name): session = boto3.Session(profile_name=sso_profile) ec2 = session.client('ec2') regions = ec2.describe_regions() return [region['RegionName'] for region in regions['Regions']] def check_guardduty_enabled(region): client = get_client('guardduty', region) detectors = client.list_detectors() return len(detectors['DetectorIds']) > 0 def ensure_eventbridge_rule(region): client = get_client('events', region) rule_name = "GuardDutyAlertRule" try: client.describe_rule(Name=rule_name) print(f"EventBridge Rule '{rule_name}' already exists in {region}.") except client.exceptions.ResourceNotFoundException: print(f"Creating EventBridge Rule '{rule_name}' in {region}...") client.put_rule( Name=rule_name, EventPattern=json.dumps({ "source": ["aws.guardduty"], "detail-type": ["GuardDuty Finding"] }) ) def ensure_sns_topic_and_subscribe_lambda(region, lambda_arn): sns_client = get_client('sns', region) topic_name = "GuardDutyAlert" topics = sns_client.list_topics() topic_arn = next((t['TopicArn'] for t in topics['Topics'] if topic_name in t['TopicArn']), None) if not topic_arn: print(f"Creating SNS Topic '{topic_name}' in {region}...") topic_arn = sns_client.create_topic(Name=topic_name)['TopicArn'] # Check current subscriptions to ensure the Lambda is subscribed subscriptions = sns_client.list_subscriptions_by_topic(TopicArn=topic_arn) subscribed = any(s['Endpoint'] == lambda_arn for s in subscriptions['Subscriptions']) if not subscribed: print(f"Subscribing Lambda to SNS Topic in {region}...") response = sns_client.subscribe( TopicArn=topic_arn, Protocol='lambda', Endpoint=lambda_arn ) subscription_arn = response['SubscriptionArn'] print(f"Subscribed with ARN: {subscription_arn}") else: print(f"Lambda already subscribed to SNS Topic in {region}.") return topic_arn def ensure_eventbridge_rule_targets(region, rule_name, target_arn): client = get_client('events', region) try: # Get current targets response = client.list_targets_by_rule(Rule=rule_name) if any(target['Arn'] == target_arn for target in response['Targets']): print(f"Target for rule '{rule_name}' is already set correctly in {region}.") else: print(f"Setting target for rule '{rule_name}' in {region}...") client.put_targets( Rule=rule_name, Targets=[ { 'Id': '1', # Identifier can be any unique string 'Arn': target_arn } ] ) except client.exceptions.ResourceNotFoundException: print(f"Rule '{rule_name}' does not exist in {region}. No targets set.") except Exception as e: print(f"Failed to set targets for rule '{rule_name}' in {region}: {str(e)}") def check_and_upload_lambda_function(region, zip_file_path): lambda_client = get_client('lambda', region) function_name = "security_alert" try: # Check if the Lambda function exists lambda_client.get_function(FunctionName=function_name) print(f"Lambda function '{function_name}' exists in {region}. Updating code...") except lambda_client.exceptions.ResourceNotFoundException: print(f"Lambda function '{function_name}' does not exist in {region}. Creating function...") # Create the Lambda function if it does not exist with open(zip_file_path, 'rb') as f: zip_content = f.read() lambda_client.create_function( FunctionName=function_name, Runtime='provided.al2023', Role='arn:aws:iam::380838443567:role/service-role/security_alert-role-d8r0k7ec', Handler='main', Code={'ZipFile': zip_content}, Environment={ 'Variables': { 'ALERT_LEVEL': '3.0', 'WEBHOOK_KEY': 'f89c38e4-1dd1-4aaf-97d0-3bd07852dd0f' } } ) def check_and_set_lambda_trigger(region, lambda_arn, topic_arn): lambda_client = get_client('lambda', region) # 检查 Lambda 函数的策略以确定是否已经允许 SNS 主题调用此函数 try: policy = lambda_client.get_policy(FunctionName=lambda_arn) policy_dict = json.loads(policy['Policy']) statements = policy_dict.get('Statement', []) # 检查是否存在允许 SNS 调用此 Lambda 函数的策略 is_allowed = any( stmt for stmt in statements if stmt.get('Action') == 'lambda:InvokeFunction' and stmt.get('Resource') == lambda_arn and 'sns.amazonaws.com' in stmt.get('Principal', {}).get('Service', '') and stmt.get('Condition', {}).get('ArnLike', {}).get('AWS:SourceArn') == topic_arn ) if not is_allowed: print(f"No permission for SNS to trigger Lambda in {region}, adding permission...") add_invoke_permission(lambda_client, lambda_arn, topic_arn) else: print(f"Permission for SNS to trigger Lambda already exists in {region}.") except lambda_client.exceptions.ResourceNotFoundException: print(f"Lambda policy not found in {region}, adding permission...") add_invoke_permission(lambda_client, lambda_arn, topic_arn) def add_invoke_permission(lambda_client, lambda_arn, topic_arn): lambda_client.add_permission( FunctionName=lambda_arn, StatementId=f'SNSInvoke-{topic_arn[-8:]}', # 使用 SNS 主题 ARN 的一部分作为唯一标识 Action='lambda:InvokeFunction', Principal='sns.amazonaws.com', SourceArn=topic_arn ) print(f"Added permission for SNS to invoke Lambda function: {lambda_arn}") def main(): regions = get_all_regions('guardduty') zip_file_path = 'main.zip' # Assuming the zip file is in the same directory as the script for region in regions: print(f"Checking GuardDuty in {region}") if check_guardduty_enabled(region): lambda_arn = f"arn:aws:lambda:{region}:380838443567:function:security_alert" ensure_eventbridge_rule(region) sns_arn = ensure_sns_topic_and_subscribe_lambda(region, lambda_arn) ensure_eventbridge_rule_targets(region, "GuardDutyAlertRule", sns_arn) check_and_upload_lambda_function(region, zip_file_path) check_and_set_lambda_trigger(region, lambda_arn, sns_arn) else: print(f"GuardDuty not enabled in {region}. Skipping this region.") if __name__ == "__main__": main()