import boto3
import ipaddress


IP_SETS = [
    {
        "name": "gsirt-ioc-ip-set-ipv4",
        "scope": "CLOUDFRONT",
        "region": "us-east-1",
        "role_name": "dev-gsirt-ioc-ip-set-deployment-role",
        "account_id": "103233932089",
    },
    {
        "name": "gsirt-ioc-ip-set-ipv4",
        "scope": "REGIONAL",
        "region": "us-east-1",
        "role_name": "dev-gsirt-ioc-ip-set-deployment-role",
        "account_id": "103233932089",
    },
    {
        "name": "gsirt-ioc-ip-set-ipv4",
        "scope": "REGIONAL",
        "region": "us-east-1",
        "role_name": "prod-gsirt-ioc-ip-set-deployment-role",
        "account_id": "437795906767",
    },
    {
        "name": "gsirt-ioc-ip-set-ipv4",
        "scope": "CLOUDFRONT",
        "region": "us-east-1",
        "role_name": "prod-gsirt-ioc-ip-set-deployment-role",
        "account_id": "437795906767",
    },
]


def assume_source_account_role(account_id: str, role_name: str, **kwargs) -> dict:
    """
    Assume IAM role in source account.

    Args:
    account_id (str): AWS account ID
    role_name (str): IAM role name to assume
    kwargs (dict): Additional parameters for assume_role

    Returns:
    dict: Dict of AWS credentials
    """
    client = boto3.client("sts")
    assume_role_response = client.assume_role(
        RoleArn=f"arn:aws:iam::{account_id}:role/{role_name}",
        RoleSessionName="ip-set-update",
        DurationSeconds=3600,
        **kwargs
    )
    credentials = assume_role_response["Credentials"]

    return credentials


def get_ip_set_by_name(client: boto3.client, name: str, scope: str) -> dict:
    """Get an IP Set by name.

    Args:
        client (boto3.client): boto3 client for WAFv2
        name (str): Name of the IP Set
        scope (str): Scope of the IP Set (CLOUDFRONT or REGIONAL)

    Returns:
        dict: IP Set details or None if not found
    """
    response = client.list_ip_sets(Scope=scope)
    for ip_set in response["IPSets"]:
        if ip_set["Name"] == name:
            return ip_set
    return None


def main():
    with open("ip_blocks.txt", "r") as file:
        addresses = [line.strip() for line in file.readlines()]
        for address in addresses:
            # Validate IP address format
            assert ipaddress.ip_network(address), "Not a valid network!"
        for ip_set in IP_SETS:
            additional_assume_role_params = {}
            if "external_id" in ip_set:
                additional_assume_role_params['ExternalId'] = ip_set["external_id"]
            credentials = assume_source_account_role(
                ip_set["account_id"],
                ip_set["role_name"],
                **additional_assume_role_params
            )
            client = boto3.client(
                "wafv2",
                region_name=ip_set["region"],
                aws_access_key_id=credentials["AccessKeyId"],
                aws_secret_access_key=credentials["SecretAccessKey"],
                aws_session_token=credentials["SessionToken"],
            )
            try:
                matching_ip_set = get_ip_set_by_name(
                    client, ip_set["name"], ip_set["scope"]
                )
                if matching_ip_set:
                    print(f"Found {matching_ip_set}")
                    response = client.update_ip_set(
                        Name=matching_ip_set["Name"],
                        Scope=ip_set["scope"],
                        Id=matching_ip_set["Id"],
                        LockToken=matching_ip_set["LockToken"],
                        Addresses=addresses,
                    )
                    print(f"IP Set {ip_set['name']} updated")
                    print(f"Response: {response}")
                else:
                    print(f"IP Set {ip_set['name']} not found.")
            except Exception as e:
                print(f"Error: {e}")


if __name__ == "__main__":
    main()