Python script to load temporary credentials for 2FA

aws_shell.py

#!/usr/bin/env python

# EB CLI doesn't support assuming roles like the AWS CLI does. This script is a workaround that requests
# temporary credentials, then exports them to the environment for packer to use.
# It then shares the AMI to the dev / prod accounts (from the cooporate account)

import ConfigParser
from os.path import expanduser
import boto3
import boto3.session
import os
import sys
import argparse
import subprocess
import uuid
import getpass
import fcntl
import ruamel.yaml as yaml
import dateutil.parser
import pytz
from datetime import datetime, timedelta

parser = argparse.ArgumentParser()
parser.add_argument('--profile', required=True)
parser.add_argument('--shell', '-s', action='store_true')
parser.add_argument('rest', nargs=argparse.REMAINDER)
args = parser.parse_args()

config = ConfigParser.ConfigParser({
    'role_arn' : None,
    'mfa_serial' : None,
})
mfa_serial = None
role_arn = None
try:
    config.read(expanduser("~/.aws/config"))
    role_arn = config.get('profile %s' % args.profile, 'role_arn')
    mfa_serial = config.get('profile %s' % args.profile, 'mfa_serial')
except:
    pass

if mfa_serial is None:
    try:
        config.read(expanduser("~/.aws/credentials"))
        mfa_serial = config.get(args.profile, 'mfa_serial')
    except:
        print "Could not find profile %s in ~/.aws/credentials or ~/.aws/config" % args.profile
        sys.exit(1)

# See if we wrote some credentials down somewhere before going all out...
aws_env = {}
with open(expanduser('~/.aws/aws_shell_cached_creds'), 'a+b') as credfile:
    credfile.seek(0)
    fcntl.flock(credfile, fcntl.LOCK_EX)
    existing_creds = None
    existing_creds_struct = {}
    try:
        _s = yaml.safe_load(credfile.read())
        if _s:
            existing_creds_struct = _s
            ts = dateutil.parser.parse(existing_creds_struct[args.profile]['AWS_CREDS_EXPIRE_AT'])
            if ts - datetime.utcnow().replace(tzinfo=pytz.UTC) > timedelta(minutes=10):
                existing_creds = existing_creds_struct[args.profile]
    except:
        pass


    if existing_creds is not None:
        aws_env = existing_creds
    else:
        aws_session = boto3.session.Session(profile_name=args.profile)
        sts = aws_session.client('sts', region_name='ap-southeast-2')

        credentials = None
        if role_arn is not None:
            asssumed_role = sts.assume_role(RoleArn=role_arn, DurationSeconds=3600, RoleSessionName='superadmin@%s' % args.profile)
            credentials = asssumed_role['Credentials']
        elif mfa_serial is not None:
            mfa_token_code = getpass.getpass('Enter MFA code: ')
            res = sts.get_session_token(DurationSeconds=3600, SerialNumber=mfa_serial, TokenCode=mfa_token_code)
            credentials = res['Credentials']
        else:
            print 'Profile needs to have either a role_arn or a mfa_serial'
            sys.exit(1)

        aws_env = {}
        aws_env['AWS_ACCESS_KEY_ID'] = credentials['AccessKeyId']
        aws_env['AWS_SECRET_ACCESS_KEY'] = credentials['SecretAccessKey']
        aws_env['AWS_SESSION_TOKEN'] = credentials['SessionToken']
        aws_env['AWS_SECURITY_TOKEN'] = credentials['SessionToken']
        aws_env['AWS_REGION'] = 'ap-southeast-2'
        aws_env['AWS_PROFILE'] = args.profile
        aws_env['AWS_CREDS_EXPIRE_AT'] = credentials['Expiration'].isoformat()

        existing_creds_struct[args.profile] = aws_env
        credfile.seek(0)
        credfile.truncate(0)
        credfile.write(yaml.dump(existing_creds_struct))

if args.shell:
    print "export AWS_ACCESS_KEY_ID=%s" % aws_env['AWS_ACCESS_KEY_ID']
    print "export AWS_SECRET_ACCESS_KEY=%s" % aws_env['AWS_SECRET_ACCESS_KEY']
    print "export AWS_SESSION_TOKEN=%s" % aws_env['AWS_SESSION_TOKEN']
    print "export AWS_SECURITY_TOKEN=%s" % aws_env['AWS_SECURITY_TOKEN']
    print "export AWS_REGION=%s" % aws_env['AWS_REGION']
    print "export AWS_PROFILE=%s" % aws_env['AWS_PROFILE']
else:
    print 'Temporary AWS crendentials have been loaded into your environment'
    print 'Spawning shell'
    new_env = {}
    new_env.update(os.environ)
    new_env.update(aws_env)
    os.execvpe(os.environ['SHELL'], [os.environ['SHELL']] + args.rest, new_env)