#!/bin/sh
set -eu

# This script runs the AWS assume-role command, captures the output, sets the
# environment variables (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY and
# AWS_SESSION_TOKEN) and executes the command given.

usage() {
    echo "$(basename "$0"): [-h|--help] ROLE_ARN|ROLE_NAME COMMAND [PARAMETER [PARAMETER ...]]"
}

command -v aws > /dev/null || { echo 'Cannot find the AWS CLI, exiting.' >&2; exit 1; }

if [ "$#" -lt 2 ]
then
    usage
    exit 1
fi

role="$1"
shift

if mfa_count="$(aws iam list-mfa-devices --query 'length(MFADevices)' 2>/dev/null)" && [ "$mfa_count" -gt 0 ]
then
    printf "Enter the MFA token: "
    read -r mfa_token
    for mfa_dev in $(aws iam list-mfa-devices --query 'MFADevices[].SerialNumber' --output text)
    do
        credentials="$(aws sts get-session-token --output text --token-code "$mfa_token" --serial-number "$mfa_dev" 2>/dev/null)" || continue
        break
    done
    if [ -z "$credentials" ]
    then
        echo 'Failed to get a temporary token.' >&2
        exit 1
    fi

    AWS_ACCESS_KEY_ID="$(echo "$credentials" | awk '{print $2}')"
    AWS_SECRET_ACCESS_KEY="$(echo "$credentials" | awk '{print $4}')"
    AWS_SESSION_TOKEN="$(echo "$credentials" | awk '{print $5}')"

    export AWS_ACCESS_KEY_ID
    export AWS_SECRET_ACCESS_KEY
    export AWS_SESSION_TOKEN

    unset AWS_SECURITY_TOKEN
fi


if [ "$role" = "${role##arn:}" ]
then
    role_arn="$(aws iam list-roles --query "Roles[?RoleName==\`${role}\`].Arn" --output text)"
else
    role_arn="$role"
fi

credentials="$(aws sts assume-role \
    --output text \
    --duration-seconds 3600 \
    --role-arn "$role_arn" \
    --role-session-name 'assume-role-cli')"

AWS_ACCESS_KEY_ID="$(echo "$credentials" | awk 'NR == 2 {print $2}')"
AWS_SECRET_ACCESS_KEY="$(echo "$credentials" | awk 'NR == 2 {print $4}')"
AWS_SESSION_TOKEN="$(echo "$credentials" | awk 'NR == 2 {print $5}')"

export AWS_ACCESS_KEY_ID
export AWS_SECRET_ACCESS_KEY
export AWS_SESSION_TOKEN

unset AWS_SECURITY_TOKEN

exec "$@"