#!/bin/bash

set -eo pipefail

# Try to discover it from the CLI config
if [ -z "$1" ]; then
    # Check that the profile exists
    grep "profile $AWS_PROFILE" ~/.aws/config > /dev/null 2>&1

    ACCOUNT=$(grep -A3 "profile $AWS_PROFILE" ~/.aws/config | grep ^account | cut -d" " -f3)
    ROLE=$(grep -A3 "profile $AWS_PROFILE" ~/.aws/config | grep ^role | cut -d" " -f3)
    ROLE_ARN="arn:aws:iam::$ACCOUNT:role/$ROLE"
else
    ROLE="$1"
fi

unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN

echo "Assuming '$ROLE_ARN' as '$USER' with profile '${AWS_PROFILE:-default}'..." >&2
creds=( $(aws sts assume-role --role-arn $ROLE_ARN --role-session-name $USER --query 'Credentials.[AccessKeyId,SecretAccessKey,SessionToken]' | tr -d ',') )

echo export AWS_ACCESS_KEY_ID="${creds[1]}"
echo export AWS_SECRET_ACCESS_KEY="${creds[2]}"
echo export AWS_SESSION_TOKEN="${creds[3]}"