import boto3
import botocore.exceptions as exceptions
import re
import os
import json

DEFAULT_PROFILE = "orcd-prod"
DEFAULT_REGION = "us-east-1"
OUTPUT_FILE_EIPS = "eips.json"
OUTPUT_FILE_VPCS = "vpcs.json"



def get_aws_accounts(ssm_client):
    aws_accounts = []

    for page in ssm_client.get_paginator('describe_parameters').paginate(
            Shared=True,
            ParameterFilters=[
                {
                    'Key': 'Name',
                    'Option': 'BeginsWith',
                    'Values': [
                        '/shared/aws-account-ids/',
                    ]
                },
            ]
    ):
        get_parameters_response = ssm_client.get_parameters(
            Names=[param['ARN'] for param in page['Parameters']],
        )
        aws_accounts.extend([json.loads(parameter['Value']) for parameter in get_parameters_response['Parameters']])

    return aws_accounts



def get_mfa_token() -> str:
    token_pattern = re.compile("^[0-9]{6}$")
    print("Enter MFA token: ")
    while True:
        mfa_token = input()
        if token_pattern.match(mfa_token):
            return mfa_token
        else:
            print("Please enter a valid MFA token: ")


def get_mfa_arn(sts, iam) -> str:
    username = sts.get_caller_identity()["Arn"].rsplit("/", 1)[1]
    mfa_devices = iam.list_mfa_devices(UserName=username)["MFADevices"]
    for mfa_d in mfa_devices:
        if mfa_d["SerialNumber"].startswith("arn:"):
            return mfa_d["SerialNumber"]


def get_mfa_creds():
    init_session = boto3.Session(
        profile_name=DEFAULT_PROFILE, region_name=DEFAULT_REGION
    )
    init_sts = init_session.client("sts")
    init_iam = init_session.client("iam")
    mfa_arn = get_mfa_arn(init_sts, init_iam)
    mfa_token = get_mfa_token()
    credentials = init_sts.get_session_token(SerialNumber=mfa_arn, TokenCode=mfa_token)[
        "Credentials"
    ]
    return {
        "aws_access_key_id": credentials["AccessKeyId"],
        "aws_secret_access_key": credentials["SecretAccessKey"],
        "aws_session_token": credentials["SessionToken"],
    }


if __name__ == "__main__":
    if os.getenv("AWS_ACCESS_KEY_ID") is None:
        creds = get_mfa_creds()
    else:
        creds = {}
    sts = boto3.client("sts", **creds)
    # Username is required because of constraints on the session name
    username = sts.get_caller_identity()["Arn"].rsplit("/", 1)[1]
    # List all regions
    client = boto3.client("ec2", **creds)
    ssm_client = boto3.client('ssm', **creds)
    regions = [region["RegionName"] for region in client.describe_regions()["Regions"]]

    ip_addresses = {}
    vpcs = {}

    accounts = [acc['account_id'] for acc in get_aws_accounts(ssm_client)]

    for acc in accounts:
        for region in regions:
            if acc != "437795906767":
                role_creds = sts.assume_role(
                    RoleArn=f"arn:aws:iam::{acc}:role/admin",
                    RoleSessionName=username,
                )["Credentials"]
                assumed_session = boto3.Session(
                    aws_access_key_id=role_creds["AccessKeyId"],
                    aws_secret_access_key=role_creds["SecretAccessKey"],
                    aws_session_token=role_creds["SessionToken"],
                    region_name=region,
                )
                ec2_client = assumed_session.client("ec2")
            else:
                ec2_client = boto3.client("ec2", **creds, region_name=region)
            try:
                response = ec2_client.describe_addresses()
                if acc not in ip_addresses:
                    ip_addresses[acc] = {}
                for address in response["Addresses"]:
                    if region not in ip_addresses[acc]:
                        ip_addresses[acc][region] = []
                    ip_addresses[acc][region].append(address["PublicIp"])

                response = client.describe_vpcs()
                if acc not in vpcs:
                    vpcs[acc] = {}
                for vpc in response['Vpcs']:
                    if region not in vpcs[acc]:
                        vpcs[acc][region] = []
                    vpcs[acc][region].append(vpc['CidrBlock'])
            except exceptions.ClientError as error:
                if error.response["Error"]["Code"] != "UnauthorizedOperation":
                    raise
                else:
                    pass
                import json

    with open(OUTPUT_FILE_EIPS, "w", encoding="utf-8") as f:
        json.dump(ip_addresses, f, ensure_ascii=False, indent=4)

    with open(OUTPUT_FILE_VPCS, "w", encoding="utf-8") as f:
        json.dump(vpcs, f, ensure_ascii=False, indent=4)