import boto3
from botocore.exceptions import ClientError
import time

# A script to enable AWS Inspector (v2) for EC2 and ECR scanning using Boto3.

# --- Configuration ---
REGION = "us-east-1"
ACCOUNT_ID = boto3.client('sts').get_caller_identity()['Account']

inspector2_client = boto3.client('inspector2', region_name=REGION)

def enable_inspector():
    """Enables Inspector for the current AWS account."""
    print("--- Enabling Amazon Inspector for EC2 and ECR ---")
    try:
        inspector2_client.enable(
            resourceTypes=['EC2', 'ECR'],
            accountIds=[ACCOUNT_ID]
        )
        print("Amazon Inspector enablement initiated.")
        # Give it a moment to update status
        time.sleep(10)
    except ClientError as e:
        if e.response['Error']['Code'] == 'ResourceAlreadyExistsException':
            print("Amazon Inspector is already enabled for this account. Skipping enablement.")
        else:
            print(f"Error enabling Inspector: {e}")
            raise

def get_inspector_status():
    """Gets and displays the status of Inspector activation."""
    print("\n--- Amazon Inspector Status ---")
    try:
        response = inspector2_client.batch_get_account_status(
            accountIds=[ACCOUNT_ID]
        )
        if response['accounts']:
            status = response['accounts'][0]['state']['status']
            print(f"Inspector status for account {ACCOUNT_ID}: {status}")
            if status == 'ENABLED':
                print("Amazon Inspector is now enabled for EC2 and ECR scanning in your account.")
                print("It may take some time for initial scans to complete and findings to appear.")
            else:
                print("Inspector is not yet fully enabled or is in an unexpected state.")
        else:
            print("Could not retrieve Inspector status for the account.")
    except ClientError as e:
        print(f"Error getting Inspector status: {e}")
        raise

def disable_inspector():
    """Disables Inspector for the current AWS account."""
    print("\n--- Disabling Amazon Inspector ---")
    try:
        inspector2_client.disable(
            resourceTypes=['EC2', 'ECR'],
            accountIds=[ACCOUNT_ID]
        )
        print("Amazon Inspector disabled.")
    except ClientError as e:
        if e.response['Error']['Code'] == 'ResourceNotFoundException':
            print("Amazon Inspector was not enabled for this account, skipping disablement.")
        else:
            print(f"Error disabling Inspector: {e}")
            raise

def main():
    try:
        enable_inspector()
        get_inspector_status()

        input("Press Enter to disable Amazon Inspector...")

    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_inspector()
        print("\n--- Amazon Inspector demonstration and cleanup complete ---")

if __name__ == "__main__":
    main()