assume_role.py

#!/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 https://gist.github.com/mbainter/b38a4cb411c0b5c1bae6 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()
config.read(os.path.expanduser("~/.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.;"
	quit(1)

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.;"
	quit(2)

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
	quit(3)

# 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 "unset AWS_SESSION_EXPIRATION;"
	print "export AWS_PROFILE=%s;" % pipes.quote(str(account))
	print "echo Account %s is setup for use.;" % account
	quit(0)

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
	quit(4)

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
			quit(0)

# 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(
    role_arn=arn,
    role_session_name=session_name
)

# 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;"
print "unset AWS_SESSION_EXPIRATION;"

# 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.bash

con() { eval $(python path/to/assume_role.py $@) ;}

credentials

[default]
aws_access_key_id = <access key>
aws_secret_access_key = <secret key>
region = us-east-1
output = json
[foo]
account_id = <account number goes here>
role = <role name goes here>
region = us-east-1
[bar]
aws_access_key_id = <access key>
aws_secret_access_key = <secret key>
region = us-east-1

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