Skip to content

Instantly share code, notes, and snippets.

What would you like to do?
#!/usr/bin/env python
# This is a trick, to output the bash commands we need to run in shell, and just execute this script inside an eval within our shell, so it imports what we need
# Possibly tie this in with for MFA support
# Will need to durably store MFA access tokens, possibly in some other env vars
# Could also store all different keys/info in different vars, to reuse as needed (lots of env vars though, file may be better)
import os
import sys
import getpass
import time
from calendar import timegm
from datetime import datetime
from dateutil.parser import *
import pipes
# These two have to be installed via pip
from ConfigParser import SafeConfigParser
from boto.sts import STSConnection
# Stupid simple way to combine all arguments given to the script as one variable, as the account
account = " ".join(sys.argv[1:])
if account == "":
account = "default"
# Setup Config Parser to read from the default .aws/credentials file
# TODO: Add check for file existence
config = SafeConfigParser()"~/.aws/credentials"))
# Read default api keys, we need these to perform the STS Assume Role call
# Possibly add support back for OptParser to read a --profile option to change which profile is used as the "base" profile
if not config.has_section('default'):
print "echo Missing 'default' section from the credentials file.;"
if not (config.has_option('default', 'aws_access_key_id') and config.has_option('default', 'aws_secret_access_key')):
print "echo Missing 'aws_access_key_id' or 'aws_secret_access_key' from 'default' section of the credentials file.;"
aws_access_key = config.get('default', 'aws_access_key_id')
aws_secret_key = config.get('default', 'aws_secret_access_key')
# Checks for the passed in account, mostly for sanity
if not config.has_section(account):
print "echo Account '%s' does not exist in the credentials file.;" % account
# If the config file is setup to use real api keys, just set the profile, and exit
# This is how the "default" provider gets used, as we have to have one set of API keys to make the rest of this work
if config.has_option(account, 'aws_access_key_id') and config.has_option(account, 'aws_secret_access_key'):
print "unset AWS_DEFAULT_REGION;"
print "unset AWS_SESSION_NAME;"
print "unset AWS_ACCESS_KEY_ID;"
print "unset AWS_SECRET_ACCESS_KEY;"
print "unset AWS_SESSION_TOKEN;"
print "export AWS_PROFILE=%s;" % pipes.quote(str(account))
print "echo Account %s is setup for use.;" % account
if not (config.has_option(account, 'account_id') and config.has_option(account, 'role')):
print "echo Account '%s' is missing the required options, 'account_id' or 'role' in the credentials file.;" % account
account_id = config.get(account, 'account_id')
role = config.get(account, 'role')
arn = 'arn:aws:iam::%s:role/%s' % (account_id, role)
session_name = '%s-%s' % (getpass.getuser(), account)
if os.environ.get('AWS_SESSION_NAME') != None:
if session_name == os.environ.get('AWS_SESSION_NAME') and os.environ.get('AWS_SESSION_EXPIRATION') != None:
stored = timegm(time.strptime(os.environ.get('AWS_SESSION_EXPIRATION').replace('Z', 'GMT'), '%Y-%m-%dT%H:%M:%S%Z'))
now = (datetime.utcnow() - datetime(1970,1,1)).total_seconds()
if stored > now:
print "echo Account %s is setup for use.;" % account
# read/set default region, default to us-east-1 if one isn't specified
region = 'us-east-1'
if config.has_option(account, 'region'):
region = config.get(account, 'region')
sts_connection = STSConnection(aws_access_key, aws_secret_key)
assumedRoleObject = sts_connection.assume_role(
# Wipe out current values, to keep things sane
print "unset AWS_DEFAULT_REGION;"
print "unset AWS_SESSION_NAME;"
print "unset AWS_ACCESS_KEY_ID;"
print "unset AWS_SECRET_ACCESS_KEY;"
print "unset AWS_SESSION_TOKEN;"
# If you swap back and forth between accounts often, this will just create new session tokens each time you swap.
# Not the cleanest solution, but it matches AWS Console behavior, and the timeout is just an hour, so you would have to swap A LOT for this to become a problem.
# The first 4 are required by any AWS SDK compatible libs to actually make requests
print "export AWS_ACCESS_KEY_ID=%s;" % pipes.quote(str(assumedRoleObject.credentials.access_key))
print "export AWS_SECRET_ACCESS_KEY=%s;" % pipes.quote(str(assumedRoleObject.credentials.secret_key))
print "export AWS_SESSION_TOKEN=%s;" % pipes.quote(str(assumedRoleObject.credentials.session_token))
print "export AWS_DEFAULT_REGION=%s;" % pipes.quote(str(region))
# The last 2 are just for our purposes, for reusing credentials for the same account if called multiple times
print "export AWS_SESSION_NAME=%s;" % pipes.quote(str(session_name))
print "export AWS_SESSION_EXPIRATION=%s;" % pipes.quote(str(assumedRoleObject.credentials.expiration))
print "echo Account %s is setup for use.;" % account
con() { eval $(python path/to/ $@) ;}
aws_access_key_id = <access key>
aws_secret_access_key = <secret key>
region = us-east-1
output = json
account_id = <account number goes here>
role = <role name goes here>
region = us-east-1
aws_access_key_id = <access key>
aws_secret_access_key = <secret key>
region = us-east-1
Copy link

mlrobinson commented Jul 12, 2015

This setup allows you to swap between different AWS accounts with one simple command, and also caches the latest credentials so multiple calls for the same assumed account don't recreate keys needlessly (but will reissue keys if they have expired). It also respects other accounts having specific keys listed, just setting the profile name as an environment variable.

The bash function calls the python script, which does the API calls and lookup logic for our .aws/credentials file. Since we are using eval in the bash function, the results are then put directly into our working environment, making subsequent calls all setup for using our API keys.

The manual parsing of the credentials file is to work around the condition of rerunning this script, swapping from one Assumed account to another, without having to return unset options, and hoping they are parsed for us in our current executing environment.

Comments in file list some shortcommings, and possible enhancements.

Copy link

alanplatt commented Aug 11, 2016

Thanks for the useful gist, I've refactored this and added MFA support plus compatibilty with the AWS CLI configs. If MFA works with the AWS CLI it will work with the new script. See aws_assume_role

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment