import pytest
import boto3
import json
from botocore.client import BaseClient
from typing import Dict
import time


# Function to get list of accounts and step functions
def get_aws_profiles_step_functions_mapping() -> Dict:
  return {
    "prod": "arn:aws:states:us-east-1:437795906767:stateMachine:prod-aws-risk-creds-exposed",
    "accounting-qa": "arn:aws:states:us-east-1:989790945997:stateMachine:qa-aws-risk-creds-exposed",
    "accounting-prod": "arn:aws:states:us-east-1:375914681009:stateMachine:prod-aws-risk-creds-exposed",
    "backup": "arn:aws:states:us-east-1:079637511089:stateMachine:backup-aws-risk-creds-exposed",
    "business-intelligence-prod": "arn:aws:states:us-east-1:313114233826:stateMachine:prod-aws-risk-creds-exposed",
    "business-solutions-prod": "arn:aws:states:us-east-1:009424872958:stateMachine:prod-aws-risk-creds-exposed",
    "business-solutions-qa": "arn:aws:states:us-east-1:814622134907:stateMachine:qa-aws-risk-creds-exposed",
    "dev": "arn:aws:states:us-east-1:103233932089:stateMachine:dev-aws-risk-creds-exposed",
    "disaster-recovery-prod": "arn:aws:states:us-east-1:682033496752:stateMachine:prod-aws-risk-creds-exposed",
    "fansifter-prod": "arn:aws:states:us-east-1:737325105821:stateMachine:prod-aws-risk-creds-exposed",
    "fansifter-qa": "arn:aws:states:us-east-1:285943604611:stateMachine:qa-aws-risk-creds-exposed",
    "gdb-apollo-dev": "arn:aws:states:us-east-1:725764004555:stateMachine:dev-aws-risk-creds-exposed",
    "gdb-apollo-prod": "arn:aws:states:us-east-1:801577982711:stateMachine:prod-aws-risk-creds-exposed",
    "gdb-atlas-dev": "arn:aws:states:us-east-1:582412345909:stateMachine:dev-aws-risk-creds-exposed",
    "gdb-atlas-prod": "arn:aws:states:us-east-1:885200931856:stateMachine:prod-aws-risk-creds-exposed",
    "gdb-crm-prod": "arn:aws:states:us-east-1:703740486246:stateMachine:dev-aws-risk-creds-exposed",
    "gdb-databricks-dev": "arn:aws:states:us-east-1:745243859288:stateMachine:dev-aws-risk-creds-exposed",
    "gdb-databricks-prod": "arn:aws:states:us-east-1:807702934854:stateMachine:prod-aws-risk-creds-exposed",
    "gdb-delphi-dev": "arn:aws:states:us-east-1:475275892927:stateMachine:dev-aws-risk-creds-exposed",
    "gdb-delphi-prod": "arn:aws:states:us-east-1:323555055331:stateMachine:prod-aws-risk-creds-exposed",
    "gdb-infra-dev": "arn:aws:states:us-east-1:483193324480:stateMachine:dev-aws-risk-creds-exposed",
    "gdb-infra-prod": "arn:aws:states:us-east-1:890208661333:stateMachine:prod-aws-risk-creds-exposed",
    "gdb-rti-dev": "arn:aws:states:us-east-1:714609459574:stateMachine:dev-aws-risk-creds-exposed",
    "gdb-rti-prod": "arn:aws:states:us-east-1:923763343190:stateMachine:prod-aws-risk-creds-exposed",
    "gdb-whitelist-prod": "arn:aws:states:eu-west-1:322697034778:stateMachine:prod-aws-risk-creds-exposed",
    "mobile-code-push-prod": "arn:aws:states:us-east-1:588827373449:stateMachine:prod-aws-risk-creds-exposed",
    "mobile-code-push-qa": "arn:aws:states:us-east-1:635057242373:stateMachine:qa-aws-risk-creds-exposed",
    "networking": "arn:aws:states:us-east-1:734458723522:stateMachine:networking-aws-risk-creds-exposed",
    "permissions-platform-prod": "arn:aws:states:us-east-1:031099521156:stateMachine:prod-aws-risk-creds-exposed",
    "permissions-platform-qa": "arn:aws:states:us-east-1:591204808501:stateMachine:qa-aws-risk-creds-exposed",
    "royaltyshare": "arn:aws:states:us-east-1:290365479392:stateMachine:dev-aws-risk-creds-exposed",
    "songwhip-prod": "arn:aws:states:us-east-1:926734670777:stateMachine:prod-aws-risk-creds-exposed",
    "songwhip-qa": "arn:aws:states:us-east-1:619719722105:stateMachine:qa-aws-risk-creds-exposed",
    "shared": "arn:aws:states:us-east-1:086679231553:stateMachine:shared-aws-risk-creds-exposed",
    "smu-dev": "arn:aws:states:us-east-1:724514206585:stateMachine:dev-aws-risk-creds-exposed",
    "smu-prod": "arn:aws:states:us-east-1:809878276326:stateMachine:prod-aws-risk-creds-exposed",
    "supply-chain-prod": "arn:aws:states:us-east-1:710271914708:stateMachine:prod-aws-risk-creds-exposed",
    "supply-chain-qa": "arn:aws:states:us-east-1:311141540202:stateMachine:qa-aws-risk-creds-exposed",
    "test-automation-qa": "arn:aws:states:us-east-1:699475937692:stateMachine:qa-aws-risk-creds-exposed",
  }


# Function to create IAM user and get Access Key
def create_user_and_get_key(iam_client: BaseClient) -> str:
  try:
    iam_client.create_user(
      UserName='test-user',
    )
  except iam_client.exceptions.EntityAlreadyExistsException:
    print("User 'test-user' already exists. Skipping user creation...")

  try:
    response: Dict = iam_client.create_access_key(
      UserName='test-user',
    )
  except iam_client.exceptions.LimitExceededException:
    # highly unlikely to happen but still better to stop here
    raise Exception("You have reached a limit of access keys for the user")

  access_key_id: str = response['AccessKey']['AccessKeyId']

  return access_key_id


# Function to delete the Access Key for the user
def delete_access_key(iam_client: BaseClient, access_key_id: str) -> None:
  iam_client.delete_access_key(
    UserName='test-user',
    AccessKeyId=access_key_id,
  )


# Function to delete the user
def delete_user(iam_client: BaseClient) -> None:
  iam_client.delete_user(
    UserName='test-user',
  )


# Function to invoke step function with payload
def invoke_step_function(sfn_client: BaseClient, sts_client: BaseClient, account_id: str, access_key_id: str, step_function_arn: str) -> str:
  payload = {
    "version": "0",
    "id": "121345678-1234-1234-1234-123456789012",
    "detail-type": "AWS Health Event",
    "source": "aws.health",
    "account": account_id,
    "time": "2016-06-05T06:27:57Z",
    "region": "us-east-1",
    "resources": [],
    "detail": {
      "eventArn": "arn:aws:health:us-east-1::event/AWS_RISK_CREDENTIALS_EXPOSED_XXXXXXXXXXXXXXXXX",
      "service": "RISK",
      "eventTypeCode": "AWS_RISK_CREDENTIALS_EXPOSED",
      "eventTypeCategory": "issue",
      "startTime": "Sat, 05 Jun 2016 15:10:09 GMT",
      "eventDescription": [
        {
          "language": "en_US",
          "latestDescription": "A description of the event will be provided here"
        }
      ],
      "affectedEntities": [
        {
          "entityValue": access_key_id
        }
      ]
    }
  }

  response: Dict = sfn_client.start_execution(
    stateMachineArn=step_function_arn,
    input=json.dumps(payload)
  )

  return response['executionArn']


def wait_for_step_function(sfn_client: BaseClient, execution_arn: str) -> None:
  while True:
    response: Dict = sfn_client.describe_execution(
      executionArn=execution_arn
    )
    status: str = response['status']

    # we might want to investigate failed, timed_out and aborted statuses
    print(f'{execution_arn} - {status}')

    if status == 'SUCCEEDED' or status == 'FAILED' or status == 'TIMED_OUT' or status == 'ABORTED':
      break

    # wait for 5 seconds before checking the status again
    time.sleep(5)


# Function to check if Access Key is deactivated
def check_key_status(iam_client: BaseClient, access_key_id: str) -> str:
  response: Dict = iam_client.list_access_keys(
    UserName='test-user',
  )

  for key in response['AccessKeyMetadata']:
    if key['AccessKeyId'] == access_key_id:
      return key['Status']


def test_key_deactivated() -> None:
  for profile_name, step_function_arn in get_aws_profiles_step_functions_mapping().items():
    region = step_function_arn.split(":")[3]
    session = boto3.Session(profile_name=profile_name, region_name=region)
    account_id = session.client('sts').get_caller_identity().get('Account')

    iam_client: BaseClient = session.client('iam')
    sfn_client: BaseClient = session.client('stepfunctions')
    sts_client: BaseClient = session.client('sts')


    access_key_id: str = create_user_and_get_key(iam_client)
    execution_arn: str = invoke_step_function(sfn_client, sts_client, account_id, access_key_id, step_function_arn)

    # Wait for step function to complete
    wait_for_step_function(sfn_client, execution_arn)

    key_status: str = check_key_status(iam_client, access_key_id)

    # clean up after the test
    delete_access_key(iam_client, access_key_id)
    delete_user(iam_client)

    assert key_status == 'Inactive', f'Access key {access_key_id} is still Active'