Rotating through AWS Profiles to Overwrite the Default

rotation.py

#!/usr/bin/env python3

"""
AWS Profile Rotation Script
"""

import argparse
import configparser
import pathlib
from copy import deepcopy
from typing import Dict, Tuple

aws_dir = pathlib.Path.home().joinpath(".aws").resolve()
config_file = aws_dir.joinpath("config")


def get_parser() -> configparser.ConfigParser:
    """
    Get the parser object
    """
    config_parser = configparser.ConfigParser()
    with config_file.open(mode="r") as config:
        config_parser.read_file(config)
    return config_parser


def get_config_section(
        parser: configparser.ConfigParser, profile: str
) -> Tuple[str, Dict[str, str]]:
    """
    Retrieve a desired section from the config

    Parameters
    ----------
    parser: configparser.ConfigParser
    profile: str

    Returns
    -------
    Tuple[str, Dict[str, str]]
    """
    config_mapping = {
        section: dict(parser.items(section=section)) for section in parser.sections()
    }
    profile_names_cleaned = [
        section_name.replace("profile ", "")
        for section_name in config_mapping.keys()
        if section_name != "profile default"
    ]
    profile_name = f"profile {profile}"
    if profile_name == "profile default":
        raise ValueError(
            "Default Profile Rotation Isn't Supported. "
            f"Available Profiles include {' | '.join(profile_names_cleaned)}"
        )
    if profile_name not in config_mapping:
        raise KeyError(
            f"That profile doesn't exist in the {config_file} file. "
            f"Available Profiles include {' | '.join(profile_names_cleaned)}"
        )
    return profile_name, config_mapping[profile_name]


def set_parser_section(
        parser: configparser.ConfigParser, options: Dict[str, str]
) -> configparser.ConfigParser:
    """
    Set Values on the parser object

    Parameters
    ----------
    parser: configparser.ConfigParser
    options: Dict[str, str]

    Returns
    -------

    """
    new_parser = deepcopy(parser)
    default_profile_name = "profile default"
    if default_profile_name not in new_parser.sections():
        new_parser.add_section(section=default_profile_name)
    for key, value in options.items():
        new_parser.set(section="profile default", option=key, value=value)
    return new_parser


def write_parser_to_file(parser: configparser.ConfigParser) -> pathlib.Path:
    """
    Write the contents of a ConfigParser to the AWS File

    Parameters
    ----------
    parser: configparser.ConfigParser

    Returns
    -------
    pathlib.Path
    """
    with config_file.open(mode="w") as config:
        parser.write(config)


if __name__ == "__main__":
    arg_parser = argparse.ArgumentParser(description="Rotate a listed AWS Profile to your default profile")
    arg_parser.add_argument("profile")
    args, _ = arg_parser.parse_known_args()
    parser = get_parser()
    print(f"Rotating to Profile: {args.profile}")
    _, options = get_config_section(parser=parser, profile=args.profile)
    parser = set_parser_section(parser=parser, options=options)
    write_parser_to_file(parser=parser)
    print(f"Rotation Complete: {args.profile}")

This gist sets up all profiles via aws-sso-util

Shell Alias:

function rotate() {
  SSO_LOGIN_ARG=${1:-${AWS_PROFILE:-Default-Profile-Name}}
  {
    aws sts get-caller-identity >/dev/null
  } || {
    aws sso login --profile ${SSO_LOGIN_ARG}
  }
  python ~/git/demo/rotate-demo/rotation.py ${SSO_LOGIN_ARG}
}