AWS Assume Role interactive utility - stores temporary session tokens and manages local credentials profile

awsrole.py

#!/usr/bin/env python3

import boto3
import argparse
import configparser
from os.path import expanduser
from botocore.exceptions import ClientError

def chose_profile()->str:
    session = boto3.Session()
    profiles = session.available_profiles
    for i in range(0, len(profiles)):
        print(f'[{i}] {profiles[i]}')
    profile_number = int(input('Choose a profile (Ctrl+C to exit): ').strip())
    return profiles[profile_number]

def get_current_account(session: boto3.Session)->int:
    client = session.client("sts")
    return int(client.get_caller_identity()['Account'])

def chose_role(session: boto3.Session)->str:
    client = session.client('iam')
    try:
        response = client.list_roles()
    except ClientError:
        print(f'Chosen profile cannot list available roles, please use -n|--assume_role_name to choose the role name to assume')
        exit(1)

    roles = []
    for role in response['Roles']:
        roles.append(role['RoleName'])
    while response.get('IsTruncated'):
        response = client.list_roles(
            Marker=response['Marker']
        )
        for role in response['Roles']:
            roles.append(role['RoleName'])
    for i in range(0, len(roles)):
        print(f'[{i}] {roles[i]}')

    role_number = int(input('Choose a role (Ctrl+C to exit): ').strip())
    return roles[role_number]

def do_assume_role(session: boto3.Session, role_arn: str, role_session_name: str, duration_seconds: int, external_id: str = False):
    sts_client = session.client("sts")
    opts = {
        'RoleArn': role_arn,
        'RoleSessionName': role_session_name,
        'DurationSeconds': duration_seconds,
    }
    if external_id:
        opts['ExternalId'] = external_id
    assumedRoleObject = sts_client.assume_role(**opts)
    return assumedRoleObject["Credentials"]

def main(args: dict):
    profile_name = args.profile
    if profile_name == "default":
        session = boto3.Session(region_name=args.region)
    elif profile_name:
        session = boto3.Session(profile_name=profile_name, region_name=args.region)
    else:
        profile_name = chose_profile()
        session = boto3.Session(profile_name=profile_name)

    external_id = args.external_id
    if external_id and not args.silent:
        print(f"WARNING: Assuming a role using ExternalID on the command line is insecure")
    if not external_id and args.prompt_external_id:
        external_id = input('Enter ExternalID (Ctrl+C to exit): ').strip()

    assume_role_name = args.assume_role_name
    if not assume_role_name:
        assume_role_name = chose_role(session)

    section_name = f'{profile_name}-{assume_role_name}'
    if args.prefix_profile:
        section_name = f'{args.prefix_profile}-{args.assume_role_name}'
    role_session_name = section_name
    if args.session_name:
        role_session_name = args.session_name

    assume_role_account = args.assume_role_account
    if not assume_role_account:
        assume_role_account = get_current_account(session)

    role_arn = f"arn:aws:iam::{assume_role_account}:role/{assume_role_name}"
    if not args.silent:
        print(f"Assuming role [{role_arn}]")
    credentials = do_assume_role(session, role_arn, role_session_name, args.duration_seconds, external_id)
    config = configparser.RawConfigParser()

    with open(args.credentials_file, "r") as f:
        config.read_file(f)
        config.remove_section(section_name)

    config.add_section(section_name)
    config.set(section_name, "aws_role_arn", role_arn)
    config.set(section_name, "region", args.region)
    config.set(section_name, "aws_access_key_id", credentials["AccessKeyId"])
    config.set(section_name, "aws_secret_access_key", credentials["SecretAccessKey"])
    config.set(section_name, "aws_session_token", credentials["SessionToken"])

    if args.default_profile:
        config.remove_section("default")
        config.add_section("default")
        config.set("default", "aws_role_arn", role_arn)
        config.set("default", "region", args.region)
        config.set("default", "aws_access_key_id", credentials["AccessKeyId"])
        config.set("default", "aws_secret_access_key", credentials["SecretAccessKey"])
        config.set("default", "aws_session_token", credentials["SessionToken"])

    with open(args.credentials_file, "w") as f:
        config.write(f)
        if not args.silent:
            print(f"\r\n\tYour access key pair has been stored in the AWS configuration file under the [{section_name}] profile.\r\n\tCredentials will expire at {'{:%Y-%m-%d %H:%M:%S}'.format(credentials['Expiration'])}\r\n\r\n\tUsage:")
            if args.default_profile:
                print("\t~$ aws sts get-caller-identity")
            else:
                print(f"\t~$ aws --profile {section_name} sts get-caller-identity")

if __name__ == '__main__':
    FILE = f"{expanduser('~')}/.aws/credentials"
    parser = argparse.ArgumentParser(description="AWS Assume Role credentials")
    parser.add_argument("-p", "--profile", help="IAM user with sts:AssumeRole (aws configure --profile <name>)")
    parser.add_argument("-a", "--assume-role-account", help="assume into aws account number. e.g. 012345678910")
    parser.add_argument("-r", "--assume-role-name", help="role name to assume")
    parser.add_argument("-i", "--external-id", help="ExternalId passed to assume")
    parser.add_argument("-I", "--prompt-external-id", help="Ask the user for the ExternalId", action='store_true')
    parser.add_argument("-s", "--duration-seconds", type=int, default=3600, help="Session duration in seconds")
    parser.add_argument("-d", "--default-profile", action="store_true")
    parser.add_argument("-q", "--silent", action='store_true', default=False)
    parser.add_argument("-P", "--prefix-profile", default='', help="prefix to use with the role name for your credential profile name")
    parser.add_argument("-R", "--region", default="ap-southeast-2")
    parser.add_argument("-N", "--session-name", default='', help="session name to use with the role_session_name parameter of sts assume role")
    parser.add_argument("-C", "--credentials-file", default=FILE, help=f"absolute path to aws credentials file (default: {FILE})")
    args = parser.parse_args()

    if args.default_profile and args.profile == "default":
        print ("can't use default profile and over write it also")
        exit(1)

    args = parser.parse_args()
    main(args)

Features

• Works cross account; use -a|--assume-role-account
• Defaults to owner account based on chosen profile
• Control session duration using -s|--duration-seconds. defaults to 3600
• Can update the default profile of the local credentials file using -d|--default-profile
• Control the profile default region, using -R|--region
• Control the session name, using -N|--session-name so CloudTrail logs can better apply non-repudiation. Defaults to profile plus role name e.g. profile-rolename
• Change the local credentials file using -C|--credentials-file
• Use a prefix for the local credentials file profile name with -P|--prefix-profile which then also becomes the default session name
• Pass as i|--external-id argument for ExternalId

It is also interactive;

• Chose a profile from the local credentials file
• Chose the Role to assume from available roles (if the chose profile has iam:ListRoles)
• Prompts you for the ExternalID is you pass the -I argument

Usage

~$ python awsrole.py -p chris -r read-only
now your temp creds are ready
~$ aws --profile chris-read-only s3 ls

Alternatively make it an executable;
~$ curl -s -o /usr/local/bin/awsrole https://gist.githubusercontent.com/chrisdlangton/0247f74619e848e7741b60a89f94c2e8/raw/4eaf2f90cb544eada0a6883eba61561224c7a49e/awsrole.py && chmod a+x /usr/local/bin/awsrole
~$ awsrole -p chris -r read-only

Ready to use;
~$ aws --profile chris-read-only s3 ls

the -d option lets you use awscli without the profile but overwrites your default profile;
~$ awsrole -p chris -r read-only -d
Now there's no need to pass in the profile
~$ aws s3 ls

and you can control the session duration in seconds;
~$ awsrole -p chris -r read-only -s 3600

Might be worth adding something to line 74 for ExternalId:
, ExternalId=externalid

and argparse of say -i

@0x646e78 good call.
that's added now, also added a prompt i used locally