generate all available AWS sso account / permission set profiles based on sso_start_url

gen_sso_profile.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Auto generate SSO profiles based on your membership."""
import argparse
from configparser import ConfigParser
import hashlib
import json
from pathlib import Path
import re
import subprocess
import sys

CACHE_DIR = Path("~/.aws/sso/cache").expanduser()


def normalize_name(name: str):
    return "-".join(
        map(str.lower, [i for i in re.split("_|-| ", name) if i.strip()])
    ).replace("-", "_")


def iter_roles(aid: str, access_token: str):
    """Iterate roles."""
    output = subprocess.check_output(
        [
            "aws",
            "sso",
            "list-account-roles",
            "--access-token",
            access_token,
            "--account-id",
            aid,
        ]
    )
    roles = json.loads(output)["roleList"]
    for obj in roles:
        print(obj)
        yield obj["roleName"]


def main(args: argparse.Namespace):
    """Run main function."""
    config = ConfigParser()
    if not args.cache_file.exists():
        print("No cache file found. Fetching...")
        print(f"Enter: {args.sso_url}")
        print(f"SSO region: us-east-2")
        print("select any account, this part doesn't matter for this script")
        subprocess.run(
            ["aws", "configure", "sso", "--profile", "sso-login"], check=True
        )
        if not args.cache_file.exists():
            raise SystemExit("Failed to fetch SSO profile.")
        print("Done generating cache file.")
    else:
        print("Cache file found. Using cache.")
    cjson = json.loads(args.cache_file.read_text())
    access_token = cjson["accessToken"]
    output = subprocess.check_output(
        ["aws", "sso", "list-accounts", "--access-token", access_token]
    )
    accounts = json.loads(output)["accountList"]
    print("FYI, if you have a lot of accounts, this might take a while.")
    for idx, obj in enumerate(accounts):
        name = normalize_name(obj["accountName"])
        aid = obj["accountId"]
        print(
            "generating profile information for account:",
            f"({idx + 1}/{len(accounts)})",
            name,
            aid,
            file=sys.stderr,
            flush=True,
        )
        for role_name in iter_roles(aid, access_token):
            profile_name = f"{name}-{role_name}"
            config[f"profile {profile_name}"] = {
                "sso_start_url": args.sso_url,
                "sso_region": "us-east-2",
                "sso_account_id": aid,
                "sso_role_name": role_name,
            }
    with args.dfile.open("w") as f:
        config.write(f)
    print(
        "Process finished. You can find your profiles in and append to ~/.aws/config",
        args.dfile,
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description=__doc__,
    )
    parser.add_argument(
        "sso_url",
        help="sso url, e.g. https://company.awsapps.com/start#/",
        type=str,
    )
    parser.add_argument(
        "destination_file",
        help="destination file",
        type=str,
    )
    args = parser.parse_args()
    args.dfile = Path(args.destination_file).expanduser()
    fhash = hashlib.sha1(args.sso_url.encode()).hexdigest()
    args.cache_file = CACHE_DIR / f"{fhash}.json"
    args.fhash = fhash
    if args.dfile.exists():
        raise SystemExit(f"{args.dfile} already exists.")
    try:
        main(args)
    except KeyboardInterrupt:
        print("Interrupted by user.", file=sys.stderr)
        sys.exit(1)