#!/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 [ "$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 "$@"
