#!/bin/bash

SESSION_FILE=~/.aws/mfa_session
PROFILE=${1:-$AWS_PROFILE}
ARGS=''

if [ "${PROFILE:0:1}" = "-" ]; then
        echo "usage: $0 [-h] [profile]"
        exit 1
fi

if [ -n "$PROFILE" ]; then
        echo "Using profile $PROFILE" >&2
        ARGS='--profile '"$PROFILE"
fi

IDENTITY_JSON=$(aws $ARGS sts get-caller-identity)
if [ $? != 0 ]; then
        if (echo "$IDENTITY_JSON" | grep -s "ExpiredToken"); then
                unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN
        else
            exit 1
        fi
fi

USER_JSON=$(aws $ARGS iam get-user)

ACCOUNT=$(echo "$IDENTITY_JSON" |grep "Account" |awk '{print $NF}' |sed -e 's/[",]*//g')
IAMUSER=$(echo "$USER_JSON" |grep 'UserName' |awk '{print $NF}' |sed -e 's/[",]*//g')

MFA_ARN="arn:aws:iam::$ACCOUNT:mfa/$IAMUSER"

echo -n "Enter MFA token for $MFA_ARN: " >&2
read MFA_TOKEN_CODE
echo ""

SESSION_JSON=$(aws $ARGS sts get-session-token --serial-number "$MFA_ARN" --token-code "$MFA_TOKEN_CODE")
if [ $? != 0 ]; then
        exit 1
fi

AWS_ACCESS_KEY_ID=$(echo "$SESSION_JSON" |grep 'AccessKeyId' | awk '{print $NF}' |sed -e 's/[",]*//g')

if [ -z "$AWS_ACCESS_KEY_ID" ]; then
        echo "Error reading AccessKeyId"
        exit 1
fi

AWS_SESSION_TOKEN="$(echo "$SESSION_JSON" |grep 'SessionToken' | awk '{print $NF}' |sed -e 's/["]*//g')"
AWS_SECRET_ACCESS_KEY=$(echo "$SESSION_JSON" | grep 'SecretAccessKey' | awk '{print $NF}' |sed -e 's/[",]*//g')

echo -e "export AWS_ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID\n"\
"export AWS_SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY\n"\
"export AWS_SESSION_TOKEN=\"$AWS_SESSION_TOKEN\"\n" |tee $SESSION_FILE

echo "Run following command: "
echo "source $SESSION_FILE"