with-aws-mfa

#!/usr/bin/env python

# this routine gets a temporary AWS key (using the MFA) to run a command

import boto3

#boto3.set_stream_logger(name='botocore')
import argparse
import os
import logging
from enum import Enum
import subprocess



class LogLevel(Enum):
    critical = 	logging.CRITICAL
    error = logging.ERROR
    warning = logging.WARNING
    info = logging.INFO
    debug = logging.DEBUG

    def __str__(self):
        return self.name

    @staticmethod
    def from_string(s):
        try:
            return LogLevel[s]
        except KeyError:
            raise ValueError()



parser = argparse.ArgumentParser()

parser.add_argument('--profile', default='default',
                    help='name of AWS profile to use')

parser.add_argument('--log_level', type=LogLevel.from_string, choices=list(LogLevel),
                    default=LogLevel.warning)

parser.add_argument('args', nargs=argparse.REMAINDER)

args = parser.parse_args()

logging.basicConfig(level=args.log_level.value)

logger=logging.getLogger(__name__)

session = boto3.Session(profile_name=args.profile)


def get_input(prompt):
    while True:
        try:
            ret = raw_input(prompt)   # Python 2.x
        except:
            ret = input(prompt)
        ret = ret.strip()
        if 0 != len(ret):
            return ret


def set_aws_keys(args):
    env = {}
    iam_client = session.resource('iam')
    sts_client = session.client('sts')
    current_user = iam_client.CurrentUser()
    user_name = current_user.user_name
    mfa_serial_number = list(current_user.mfa_devices.all())[0].serial_number
    current_mfa_value = get_input('Enter current MFA value for %s: ' % user_name)
    response = sts_client.get_session_token(SerialNumber = mfa_serial_number,
                                            TokenCode = current_mfa_value)
    # set the environment variables
    credentials = response['Credentials']
    os.environ['AWS_ACCESS_KEY_ID'] = credentials['AccessKeyId']
    os.environ['AWS_SECRET_ACCESS_KEY'] = credentials['SecretAccessKey']
    os.environ['AWS_SESSION_TOKEN'] = credentials['SessionToken']

set_aws_keys(args)
    
os.system(" ".join(args.args))