import boto3
from botocore.exceptions import ClientError
import time

# A script to enable AWS Security Hub and then clean up using Boto3.

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

securityhub_client = boto3.client('securityhub', region_name=REGION)

def enable_security_hub():
    """Enables Security Hub for the current AWS account."""
    print("--- Enabling AWS Security Hub ---")
    try:
        response = securityhub_client.enable_security_hub()
        hub_arn = response['HubArn']
        print(f"AWS Security Hub enabled. Hub ARN: {hub_arn}")
        return hub_arn
    except ClientError as e:
        if e.response['Error']['Code'] == 'ResourceConflictException':
            print("AWS Security Hub is already enabled. Fetching Hub ARN.")
            response = securityhub_client.describe_hub()
            return response['HubArn']
        else:
            print(f"Error enabling Security Hub: {e}")
            raise

def disable_security_hub():
    """Disables Security Hub for the current AWS account."""
    print("\n--- Disabling AWS Security Hub ---")
    try:
        securityhub_client.disable_security_hub()
        print("AWS Security Hub disabled.")
    except ClientError as e:
        if e.response['Error']['Code'] == 'ResourceNotFoundException':
            print("AWS Security Hub was not enabled, skipping disablement.")
        else:
            print(f"Error disabling Security Hub: {e}")
            raise

def main():
    hub_arn = None
    try:
        hub_arn = enable_security_hub()

        print("\n--- AWS Security Hub Enabled Successfully! ---")
        print(f"Hub ARN: {hub_arn}")
        print("Security Hub is now collecting security findings for your account.")

        input("Press Enter to disable AWS Security Hub...")

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

if __name__ == "__main__":
    main()