a bash / zsh compatible alias that can assist in switching between aws profiles and triggering sso login if necessary

aws-profile.sh

aws-profile() {
    read -r -d '' SCRIPT <<'EOF'
import argparse
from configparser import RawConfigParser
import os
import shutil
import subprocess
import sys
from textwrap import dedent

# interactive output needs to be over stderr to avoid being captured
# by the bash script looking for a profile name on stdout
out = lambda *a, **kw: print(*a, **kw, file=sys.stderr)

def parse_profiles():
    config_profile_prefix = 'profile '
    config_path = os.path.expanduser('~/.aws/config')

    config = RawConfigParser()
    config.read(config_path)

    profiles = []
    for section in config.sections():
        if section.startswith(config_profile_prefix):
            profile = section[len(config_profile_prefix):]
            profiles.append(profile)

    return profiles

def select_profile_manual(profiles, current_profile=None):
    out('--- Choose AWS profile ---')
    for idx, profile in enumerate(profiles, 1):
        suffix = ' ** current **' if profile == current_profile else ''
        out(f'{idx}: {profile}{suffix}')
    out()

    while True:
        out(f'Pick one: [1-{len(profiles)}] ', end='')
        choice = input().strip()
        try:
            if not choice and current_profile:
                return current_profile
            if choice == 'q':
                sys.exit(0)
            choice = int(choice) - 1
            if not 0 <= choice < len(profiles):
                raise Exception
        except Exception:
            out('Unrecognized option, try again.')
        else:
            break

    return profiles[choice]

def select_profile_fzf(profiles, current_profile=None):
    args = ['fzf', '--height', '40%']
    if current_profile is not None:
        try:
            idx = profiles.index(current_profile)
            # pull the current profile to the front of the list so that
            # it's the default selection
            profiles = [profiles[idx]] + profiles[:idx] + profiles[idx + 1 :]
        except ValueError:
            pass
    in_profiles = '\n'.join(profiles)
    p = subprocess.run(args, input=in_profiles, text=True, stdout=subprocess.PIPE)
    if p.returncode != 0:
        return
    # the selected profile should be the last line in the output
    profile = p.stdout.strip().split()[-1].strip()
    return profile

def parse_args():
    try:
        args = sys.argv[sys.argv.index('--') + 1:]
    except IndexError:
        args = sys.argv[1:]

    parser = argparse.ArgumentParser(prog='aws-profile', add_help=False)
    parser.add_argument('-h', '--help', action='store_true')
    parser.add_argument('-c', '--current', action='store_true')
    parser.add_argument('--unset', action='store_true')
    parser.add_argument('profile', nargs='?')
    args = parser.parse_args(args)

    # take over help so that we can use stderr and exit with an error
    # to avoid the output being eval'd
    if args.help:
        parser.print_help(file=sys.stderr)
        raise SystemExit(1)
    return args

def main():
    args = parse_args()

    current_profile = os.getenv('AWS_PROFILE', '')
    if args.current:
        out(f'Current AWS_PROFILE={current_profile}')
        return 1

    if args.unset:
        print(dedent(
            '''
            unset AWS_ACCESS_KEY_ID
            unset AWS_SECRET_ACCESS_KEY
            unset AWS_SESSION_TOKEN
            unset AWS_PROFILE
            unset AWS_DEFAULT_PROFILE
            '''
        ))
        return 0

    profile = args.profile
    if not profile:
        profiles = parse_profiles()
        if not profiles:
            out('No profiles found in ~/.aws/config')
            return 1
        if shutil.which('fzf') is None:
            profile = select_profile_manual(profiles, current_profile)
        else:
            profile = select_profile_fzf(profiles, current_profile)

    if profile is None:
        out(f'No profile selected.')
        return 1

    args = ['aws', 'configure', 'list']
    env = {**dict(os.environ), 'AWS_PROFILE': profile}
    p = subprocess.run(args, capture_output=True, text=True, env=env)
    
    stderr = p.stderr.lower()
    if p.returncode != 0 and (
        'error loading sso token' in stderr
        or 'error when retrieving token' in stderr
        or 'sso session associated with this profile has expired' in stderr
    ):
        out('Session is invalid, triggering an SSO login ...')
        args = ['aws', 'sso', 'login', '--no-browser']
        p = subprocess.run(args, stdout=sys.stderr, env=env, text=True)
        if p.returncode != 0:
            raise SystemExit(1)

    elif p.returncode != 0:
        out(p.stderr)
        raise SystemExit(p.returncode)

    out(f'Activated AWS_PROFILE={profile}')
    print(dedent(
        f'''
        unset AWS_ACCESS_KEY_ID
        unset AWS_SECRET_ACCESS_KEY
        unset AWS_SESSION_TOKEN
        export AWS_PROFILE={profile}
        export AWS_DEFAULT_PROFILE={profile}
        '''
    ))
    return 0

if __name__ == '__main__':
    try:
        raise SystemExit(main())
    except KeyboardInterrupt:
        raise SystemExit(1)
EOF

    result=$(python3 -c "$SCRIPT" -- "$@")
    rc=$?
    if [ $rc -ne 0 ]; then
        return $rc
    fi
    eval "$result"
}