import boto3
from botocore.exceptions import ClientError
import time

# A script to create a Customer Master Key (CMK) in AWS KMS,
# create an alias, enable key rotation, and then clean up.

# --- Configuration ---
REGION = "us-east-1"
CMK_DESCRIPTION = "My Boto3 Test CMK"
CMK_ALIAS = "alias/MyBoto3CMK"
CMK_TAG_KEY = "Project"
CMK_TAG_VALUE = "KMSDemo"

kms_client = boto3.client('kms', region_name=REGION)

def create_cmk():
    """Creates a new symmetric customer-managed CMK."""
    print("--- Creating Customer Master Key (CMK) ---")
    try:
        response = kms_client.create_key(
            Description=CMK_DESCRIPTION,
            KeyUsage='ENCRYPT_DECRYPT',
            KeySpec='SYMMETRIC_DEFAULT',
            Tags=[{'TagKey': CMK_TAG_KEY, 'TagValue': CMK_TAG_VALUE}]
        )
        cmk_id = response['KeyMetadata']['KeyId']
        print(f"CMK created with ID: {cmk_id}")
        return cmk_id
    except ClientError as e:
        print(f"Error creating CMK: {e}")
        raise

def create_alias(cmk_id):
    """Creates an alias for the CMK."""
    print(f"\n--- Creating Alias '{CMK_ALIAS}' for CMK '{cmk_id}' ---")
    try:
        kms_client.create_alias(AliasName=CMK_ALIAS, TargetKeyId=cmk_id)
        print("Alias created.")
    except ClientError as e:
        if e.response['Error']['Code'] == 'AlreadyExistsException':
            print(f"Alias '{CMK_ALIAS}' already exists. Skipping creation.")
        else:
            print(f"Error creating alias: {e}")
            raise

def enable_key_rotation(cmk_id):
    """Enables automatic key rotation for the CMK."""
    print(f"\n--- Enabling automatic key rotation for CMK '{cmk_id}' ---")
    try:
        kms_client.enable_key_rotation(KeyId=cmk_id)
        print("Key rotation enabled.")
    except ClientError as e:
        print(f"Error enabling key rotation: {e}")
        raise

def cleanup_resources(cmk_id):
    """Cleans up all created KMS resources."""
    print(f"\n--- Cleaning up resources ---")

    # Disable CMK
    print(f"Disabling CMK '{cmk_id}'...")
    try:
        kms_client.disable_key(KeyId=cmk_id)
        print("CMK disabled.")
    except ClientError as e:
        print(f"Error disabling CMK: {e}")

    # Delete Alias
    print(f"Deleting Alias '{CMK_ALIAS}'...")
    try:
        kms_client.delete_alias(AliasName=CMK_ALIAS)
        print("Alias deleted.")
    except ClientError as e:
        print(f"Error deleting alias: {e}")

    # Schedule CMK for deletion (minimum 7 days)
    print(f"Scheduling CMK '{cmk_id}' for deletion in 7 days...")
    try:
        kms_client.schedule_key_deletion(KeyId=cmk_id, PendingWindowInDays=7)
        print("CMK scheduled for deletion.")
    except ClientError as e:
        print(f"Error scheduling CMK for deletion: {e}")

def main():
    cmk_id = None
    try:
        cmk_id = create_cmk()
        create_alias(cmk_id)
        enable_key_rotation(cmk_id)

        print("\n--- KMS CMK setup complete! ---")
        print(f"CMK ID: {cmk_id}")
        print(f"CMK Alias: {CMK_ALIAS}")

        input("Press Enter to disable the CMK and schedule it for deletion...")

    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 cmk_id:
            cleanup_resources(cmk_id)
        print("\n--- KMS CMK demonstration and cleanup complete ---")

if __name__ == "__main__":
    main()