Generate AWS sso config file with profiles for all accounts in organization

generate_aws_sso_config.py

import boto3
import argparse
import os
import sys
import configparser
import logging
import shutil  # newly added import

# Constants
ROLES = ["Admin", "GodMode", "ReadOnly"]

# Configure logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)

"""
Generate AWS SSO Config

This script generates AWS SSO configuration based on the provided SSO session name.
It retrieves the SSO session name from the command-line arguments and uses it to
fetch the session region and list of accounts associated with the SSO session.
It then generates AWS profiles for each account and updates the configuration file.
Note: This script requires an existing profile for a management account in the AWS Organization.
It will handle duplicate SSO profiles by updating the existing entries.

Prerequisite:
Run this command first:
aws configure sso

Arguments:
--sso-session (str): The organization SSO session name. (required)
--config-file (str): Optional path to AWS config file. Defaults to ~/.aws/config

Example:
python generate_aws_sso_config.py --sso-session my_sso_session
"""


def list_accounts_with_sso():
    try:
        sso_client = boto3.client("organizations")
        paginator = sso_client.get_paginator("list_accounts")
        iterator = paginator.paginate(PaginationConfig={"MaxResults": 20})
        accounts = []
        for page in iterator:
            # Filter only ACTIVE accounts
            active_accounts = [
                acc for acc in page["Accounts"] if acc["Status"] == "ACTIVE"
            ]
            accounts.extend(active_accounts)
        return accounts
    except Exception as e:
        logging.error(f"Failed to list accounts: {e}")
        sys.exit(1)


def read_config(config_path):
    config = configparser.ConfigParser()
    config.read(os.path.expanduser(config_path))
    return config


def new_config(config_path):
    config = configparser.ConfigParser()
    return config


def config_session_region(config, sso_session_name):
    section = f"sso-session {sso_session_name}"
    if section not in config.sections():
        logging.error(f"SSO session {sso_session_name} not found")
        sys.exit(1)
    region = config[section]["sso_region"]
    return region


def generate_aws_profiles(config, session_name, accounts, region):
    for account in accounts:
        account_name = account["Name"]
        profile_name = account_name.replace(" ", "--")
        account_id = account["Id"]
        for role in ROLES:
            profile_section = f"profile {profile_name}-{role}"

            if profile_section in config.sections():
                logging.warning(
                    f"Duplicate profile '{profile_section}' found. Updating existing entry."
                )
                profile_section = f"profile {profile_name}--{role}"

            config.add_section(profile_section)
            config.set(profile_section, "sso_session", session_name)
            config.set(profile_section, "sso_account_id", account_id)
            config.set(profile_section, "sso_role_name", role)
            config.set(profile_section, "region", region)


def write_sso_session(config, sso_session_name, old_config):
    # This function should write the sso-session to the config
    sso_session_header = f"sso-session {sso_session_name}"
    old_section = old_config[sso_session_header]
    logging.info(f"Adding SSO session {sso_session_name} to config")
    # print all entries in the old section
    #for key in old_section.keys():
    #    logging.info(f"{key} = {old_section[key]}")
    logging.info(f"{dict(old_section)}")
    config.add_section(sso_session_header)
    config[sso_session_header] = old_section


def write_config(config, config_path):
    expanded_path = os.path.expanduser(config_path)
    if os.path.exists(expanded_path):
        backup_path = expanded_path + ".backup"
        shutil.move(expanded_path, backup_path)
        logging.info(f"Config file moved to backup at {backup_path}")
    with open(expanded_path, "w") as config_file:
        config.write(config_file)


def main():
    parser = argparse.ArgumentParser(description="Generate AWS SSO Config")
    parser.add_argument(
        "--sso-session", required=True, help="The organization SSO session name"
    )
    parser.add_argument(
        "--config-file", default="~/.aws/config", help="Path to AWS config file"
    )
    args = parser.parse_args()
    sso_session_name = args.sso_session
    config_file = args.config_file

    old_config = read_config(config_file)
    config = new_config(config_file)
    # write the sso-session to the new config
    write_sso_session(config, sso_session_name, old_config)
    region = config_session_region(old_config, sso_session_name)
    accounts = list_accounts_with_sso()
    logging.info(f"Found {len(accounts)} accounts")
    generate_aws_profiles(config, sso_session_name, accounts, region)
    write_config(config, config_file)
    logging.info("AWS SSO configuration updated successfully.")


if __name__ == "__main__":
    if "AWS_PROFILE" not in os.environ:
        logging.error("AWS_PROFILE is not set")
        sys.exit(1)
    if "AWS_REGION" not in os.environ:
        logging.error("AWS_REGION must be set")
        sys.exit(1)
    main()