import boto3
from botocore.exceptions import ClientError
import time

# A script to enable AWS Shield Advanced and then clean up using Boto3.

# --- Configuration ---
REGION = "us-east-1"

shield_client = boto3.client('shield', region_name=REGION)

def create_shield_advanced_subscription():
    """Creates a subscription to AWS Shield Advanced."""
    print("--- Creating AWS Shield Advanced Subscription ---")
    try:
        response = shield_client.create_subscription()
        subscription_arn = response['Subscription']['SubscriptionArn']
        print(f"AWS Shield Advanced subscription created. ARN: {subscription_arn}")
        return subscription_arn
    except ClientError as e:
        if e.response['Error']['Code'] == 'ResourceAlreadyExistsException':
            print("AWS Shield Advanced subscription already exists. Fetching ARN.")
            response = shield_client.describe_subscription()
            return response['Subscription']['SubscriptionArn']
        else:
            print(f"Error creating subscription: {e}")
            raise

def delete_shield_advanced_subscription():
    """Deletes the Shield Advanced subscription."""
    print("\n--- Deleting AWS Shield Advanced subscription ---")
    try:
        shield_client.delete_subscription()
        print("AWS Shield Advanced subscription deleted.")
    except ClientError as e:
        if e.response['Error']['Code'] == 'ResourceNotFoundException':
            print("AWS Shield Advanced subscription not found, skipping deletion.")
        else:
            print(f"Error deleting subscription: {e}")
            raise

def main():
    subscription_arn = None
    try:
        subscription_arn = create_shield_advanced_subscription()

        print("\n--- AWS Shield Advanced Enabled Successfully! ---")
        print(f"Subscription ARN: {subscription_arn}")
        print("Your account is now protected by Shield Advanced.")

        input("Press Enter to delete the AWS Shield Advanced subscription...")

    except ClientError as e:
        print(f"An AWS client error occurred: {e}")
    except Exception as e:
        print(f"An unexpected error occurred: {e}")
    finally:
        delete_shield_advanced_subscription()
        print("\n--- AWS Shield Advanced demonstration and cleanup complete ---")

if __name__ == "__main__":
    main()