import boto3
from botocore.exceptions import ClientError
import time
import json

# A script to create a basic AWS WAF Web ACL with a managed rule group using Boto3.

# --- Configuration ---
REGION = "us-east-1"
WEB_ACL_NAME = "MyBoto3WafWebACL"
WEB_ACL_DESCRIPTION = "Web ACL for Boto3 demo"
MANAGED_RULE_GROUP_NAME = "AWSManagedRulesCommonRuleSet"

wafv2_client = boto3.client('wafv2', region_name=REGION)

def create_web_acl():
    """Creates a WAF Web ACL."""
    print(f"--- Creating WAF Web ACL: {WEB_ACL_NAME} ---")
    try:
        response = wafv2_client.create_web_acl(
            Name=WEB_ACL_NAME,
            Scope='REGIONAL',
            DefaultAction={'Allow': {}},
            VisibilityConfig={
                'SampledRequestsEnabled': True,
                'CloudWatchMetricsEnabled': True,
                'MetricName': f"{WEB_ACL_NAME}Metric"
            },
            Description=WEB_ACL_DESCRIPTION
        )
        web_acl_arn = response['Summary']['ARN']
        web_acl_id = response['Summary']['Id']
        print(f"Web ACL created with ARN: {web_acl_arn}")
        return web_acl_arn, web_acl_id
    except ClientError as e:
        if e.response['Error']['Code'] == 'WAFDuplicateItemException':
            print(f"Web ACL '{WEB_ACL_NAME}' already exists. Fetching details.")
            response = wafv2_client.list_web_acls(Scope='REGIONAL', Name=WEB_ACL_NAME)
            if response['WebACLs']:
                web_acl_arn = response['WebACLs'][0]['ARN']
                web_acl_id = response['WebACLs'][0]['Id']
                return web_acl_arn, web_acl_id
            else:
                raise Exception("Could not find existing Web ACL.")
        else:
            print(f"Error creating Web ACL: {e}")
            raise

def add_managed_rule_group(web_acl_id, web_acl_name):
    """Adds a managed rule group to the Web ACL."""
    print(f"\n--- Adding Managed Rule Group: {MANAGED_RULE_GROUP_NAME} to Web ACL ---")
    try:
        # Get current LockToken before updating
        response = wafv2_client.get_web_acl(Name=web_acl_name, Scope='REGIONAL', Id=web_acl_id)
        lock_token = response['LockToken']

        wafv2_client.update_web_acl(
            Name=web_acl_name,
            Scope='REGIONAL',
            Id=web_acl_id,
            LockToken=lock_token,
            DefaultAction={'Allow': {}},
            Rules=[
                {
                    'Name': 'ManagedCommonRuleSet',
                    'Priority': 1,
                    'Statement': {
                        'ManagedRuleGroupStatement': {
                            'VendorName': 'AWS',
                            'Name': MANAGED_RULE_GROUP_NAME
                        }
                    },
                    'Action': {'Block': {}},
                    'VisibilityConfig': {
                        'SampledRequestsEnabled': True,
                        'CloudWatchMetricsEnabled': True,
                        'MetricName': 'ManagedCommonRuleSetMetric'
                    }
                }
            ],
            VisibilityConfig={
                'SampledRequestsEnabled': True,
                'CloudWatchMetricsEnabled': True,
                'MetricName': f"{WEB_ACL_NAME}Metric"
            }
        )
        print("Managed Rule Group added to Web ACL.")
    except ClientError as e:
        print(f"Error adding managed rule group: {e}")
        raise

def delete_web_acl(web_acl_id, web_acl_name):
    """Deletes the WAF Web ACL."""
    print(f"\n--- Deleting WAF Web ACL: {WEB_ACL_NAME} ---")
    try:
        # Get current LockToken before deleting
        response = wafv2_client.get_web_acl(Name=web_acl_name, Scope='REGIONAL', Id=web_acl_id)
        lock_token = response['LockToken']

        wafv2_client.delete_web_acl(
            Name=web_acl_name,
            Scope='REGIONAL',
            Id=web_acl_id,
            LockToken=lock_token
        )
        print("Web ACL deleted.")
    except ClientError as e:
        if e.response['Error']['Code'] == 'WAFNonexistentItemException':
            print(f"Web ACL '{WEB_ACL_NAME}' not found, skipping deletion.")
        else:
            print(f"Error deleting Web ACL: {e}")
            raise

def main():
    web_acl_arn = None
    web_acl_id = None
    try:
        web_acl_arn, web_acl_id = create_web_acl()
        add_managed_rule_group(web_acl_id, WEB_ACL_NAME)

        print("\n--- WAF Web ACL Setup Complete! ---")
        print(f"Web ACL ARN: {web_acl_arn}")
        print("You can now associate this Web ACL with an ALB, CloudFront distribution, or API Gateway.")

        input("Press Enter to delete the WAF Web ACL...")

    except ClientError as e:
        print(f"An AWS client error occurred: {e}")
    except Exception as e:
        print(f"An unexpected error occurred: {e}")
    finally:
        if web_acl_id:
            delete_web_acl(web_acl_id, WEB_ACL_NAME)
        print("\n--- WAF Web ACL demonstration and cleanup complete ---")

if __name__ == "__main__":
    main()