#!/usr/bin/python
import boto3
import subprocess
import sys

PROXY_PROFILE = "gdb-infra-dev"

ACCOUNT_MAPPING = {
    "gdb-apollo-dev": "725764004555",
    "gdb-artistapp-dev": "714609459574",
    "gdb-core-dev": "582412345909",
    "gdb-databricks-dev": "745243859288",
    "gdb-datasci-dev": "313114233826",
    "gdb-delphi-dev": "475275892927",
    "gdb-infra-dev": "483193324480",
    "gdb-prodhub-dev": "703740486246",
    "gdb-smu-dev": "724514206585",
    "gdb-whitelist-dev": "814622134907",
    "gdb-apollo-prod": "801577982711",
    "gdb-artistapp-prod": "923763343190",
    "gdb-core-prod": "885200931856",
    "gdb-databricks-prod": "807702934854",
    "gdb-delphi-prod": "323555055331",
    "gdb-infra-prod": "890208661333",
    "gdb-smu-prod": "809878276326",
    "gdb-whitelist-prod": "009424872958",
    #"gdb-whitelist-legacy": "322697034778"
}


def update_user_key(acc: dict, username: str, use_role: bool) -> None:
    if use_role:
        acc_session = boto3.Session(
            aws_access_key_id=acc["creds"]["AccessKeyId"],
            aws_secret_access_key=acc["creds"]["SecretAccessKey"],
            aws_session_token=acc["creds"]["SessionToken"])
    else:
        acc_session = boto3.Session(profile_name=acc['alias'])
    # Check whether there are free slot for the key
    client = acc_session.client('iam')
    existing_keys = client.list_access_keys(
        UserName=username)["AccessKeyMetadata"]
    # If there are two existing keys, remove inactive keys or the old one
    if len(existing_keys) > 1:
        if existing_keys[0]["Status"] == "Inactive":
            client.delete_access_key(
                UserName=username, AccessKeyId=existing_keys[0]["AccessKeyId"])
        elif existing_keys[1]["Status"] == "Inactive":
            client.delete_access_key(
                UserName=username, AccessKeyId=existing_keys[1]["AccessKeyId"])
        else:
            client.delete_access_key(
                UserName=username, AccessKeyId=existing_keys[0]["AccessKeyId"])
    # create new key and write it to profile
    new_keys = client.create_access_key(UserName=username)["AccessKey"]
    subprocess.run([
        "aws", "configure", "--profile", f"{acc['alias']}", "set",
        "aws_access_key_id", new_keys["AccessKeyId"]
    ])
    subprocess.run([
        "aws", "configure", "--profile", f"{acc['alias']}", "set",
        "aws_secret_access_key", new_keys["SecretAccessKey"]
    ])


if __name__ == "__main__":

    #get username
    session = boto3.Session(profile_name=PROXY_PROFILE)
    sts_client = session.client('sts')
    username = sts_client.get_caller_identity()['Arn'].rsplit('/')[-1]

    accounts = []

    use_role = True if len(sys.argv) > 1 and sys.argv[1] == '--role' else False

    for alias, id in ACCOUNT_MAPPING.items():
        if use_role:
            creds = sts_client.assume_role(
                RoleArn=f"arn:aws:iam::{id}:role/cross_account_terraform_role",
                RoleSessionName=f"{username}-key-rotation")
            print("Assumed role in {0}".format(alias))
            accounts.append({
                "alias": alias,
                "id": id,
                "creds": creds['Credentials']
            })
        else:
            accounts.append({"alias": alias, "id": id})

    for acc in accounts:
        update_user_key(acc, username, use_role)