Set profile-based AWS environment variables with support for assume-role

README.md

About

This script sets AWS_DEFAULT_REGION, AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY and AWS_SESSION_TOKEN (where applicable) from the values configured in ~/.aws/config and ~/.aws/credentials This is useful for cases where the tool being used doesn't support --profile being passed in, i.e. most unit/system tests

It also can handle cases where the credentials aren't static (i.e. assume-role giving temporary credentials), it uses the awscli's implementation to resolve and cache those credentials. This is very handy for when a tool is using the aws-sdk to load credentials and can read the AWS_PROFILE environment variable but can't handle the assume-role, usuaully due to an MFA prompt. In those cases it would reolve the tokens and set them as env vars. Calling it again would refersh them if required.

Usage:

put in somewhere like ~/bin/awsenv:

cp awsenv.py ~/bin/awsenv
chmod +x ~/bin/awsenv

Then add the following wrapper function to your ~/.bashrc:

setaws () {
   [[ $# -gt 0 ]] && eval "$(~/bin/awsenv $@)"
}

to set a profile: setaws myprofile

to view current: awsenv

to clear: setaws --unset

awsenv.py

#!/usr/bin/env python3
#
# adapted from https://github.com/bdclark/awsenv/blob/master/awsenv
#
from __future__ import print_function
import os
import sys

try:
    import ConfigParser as configparser
except:
    import configparser

# Track all ENV variables that we may set so we know what to unset later"
# Key is environment variable name, value is ini setting from aws config
#
# environment variables used by awscli: 
#   https://docs.aws.amazon.com/cli/latest/topic/config-vars.html#
aws_dict = {
    'AWS_ACCESS_KEY_ID': 'aws_access_key_id',
    'AWS_SECRET_ACCESS_KEY': 'aws_secret_access_key',
    'AWS_DEFAULT_REGION': 'region',
    'AWS_SESSION_TOKEN': 'aws_session_token',
    'AWS_PROFILE': 'aws_profile',
}


def get_option(config, section, option):
    """Return value of ini option in section; empty string if not found"""
    try:
        return config.get(section, option)
    except configparser.NoOptionError:
        return ''


def show_env(config):
    """Print AWS-related environment vars, and which profile(s) they match"""
    print()
    print('Current AWS Environment Variables:')
    print()
    for env_setting, ini_opt in aws_dict.items():
        ini_value = os.environ.get(env_setting)
        if ini_value is None:
            print(env_setting, "not set")
        else:
            match_sections = []
            for section in config.sections():
                if get_option(config, section, ini_opt) == ini_value:
                    match_sections.append(section.replace('profile ', ''))
            if len(match_sections):
                print(env_setting, "matches", match_sections)
            else:
                print(env_setting, "set =", ini_value)


def unset_values():
    """return string representation of command to unset environment vars"""
    return "unset {0} AWS_PROFILE; ".format(' '.join(aws_dict.keys()))


def exit_out(message, error=False):
    out = sys.stderr if error else sys.stdout
    print(message, file=out)
    exit(1 if error else 0)


profile = ''
config_files = []

if len(sys.argv) > 1:
    if sys.argv[1] == '--unset':
        exit_out(unset_values())
    else:
        profile = sys.argv[1]

if os.environ.get('AWS_CONFIG_FILE'):
    config_files.append(os.path.expanduser(os.environ.get('AWS_CONFIG_FILE')))
else:
    config_files.append(os.path.expanduser('~/.aws/config'))

config_files.append(os.path.expanduser('~/.aws/credentials'))

if not config_files:
    exit_out("Error: no valid AWS config file(s) found", True)

configs = configparser.ConfigParser()
configs.read(config_files)

if not profile:
    show_env(configs)
    exit(0)


def create_cli_session(profile_name):
    from awscli import EnvironmentVariables
    import botocore.session
    from awscli.plugin import load_plugins

    session = botocore.session.Session(EnvironmentVariables)
    # the following loads the hooks for aws' cli for caching assume role credentials
    import argparse
    args = argparse.ArgumentParser().parse_args([])
    args.command = 'help' # dummy command
    event_handlers = session.get_component('event_emitter')
    load_plugins(session.full_config.get('plugins', {}), event_hooks=event_handlers)
    # sets the profile to load
    session.set_config_variable('profile', profile_name)
    session.emit('session-initialized', session=session, parsed_args=args)
    return session


def lookup_settings(config, profile_name):
    settings = {}
    if config.has_section('profile ' + profile_name) or config.has_section(profile):
        for env_setting, ini_opt in aws_dict.items():
            if profile != 'default':
                # Profiles in ~/.aws/config begin with "profile ", except default
                if config.has_section('profile ' + profile_name):
                    value = get_option(config, 'profile ' + profile_name, ini_opt)
                    if value:
                        settings[env_setting] = value
            # i.e. for ~/.aws/credentials (profiles are stored with no prefix)
            if config.has_section(profile_name):
                value = get_option(config, profile_name, ini_opt)
                if value:
                    settings[env_setting] = value

        # no keys configured so probably it's an assume role
        if 'AWS_ACCESS_KEY_ID' not in settings and get_option(config, 'profile ' + profile_name, 'role_arn'):
            print("Profile has no AWS_ACCESS_KEY_ID, so using awscli to determine keys", file=sys.stderr)
            session = create_cli_session(profile_name)
            credentials = session.get_credentials()
            settings['AWS_ACCESS_KEY_ID'] = credentials.access_key
            settings['AWS_SECRET_ACCESS_KEY'] = credentials.secret_key
            if credentials.token:
                settings['AWS_SESSION_TOKEN'] = credentials.token
        settings['AWS_PROFILE'] = profile

        return settings
    else:
        exit_out("Error: profile {0} not found.".format(profile_name), True)


env_settings = lookup_settings(configs, profile)

print(unset_values())
if env_settings:
    output = 'export '
    output += ' '.join("{0}={1}".format(k, v) for (k, v) in env_settings.items())
    print(output)