#!/bin/bash

# requires jq - the command-line JSON processor - to be installed

function get-session-token() {
AWSMFA_SERIAL="default_DEVICE"
AWSCLI_STS_COMMAND="get-session-token"

unset AWSCLI_OPTIONS
unset AWSCLI_STS_OPTIONS
unset AWSMFA_CREDENTIALS

OPTIND=1
while getopts ":p:r:h" opt
do
  case ${opt} in
    h)
      echo "Usage: get-session-token [-p awscli-profile] [-r AWS_IAM_ROLE_ARN] | -h"
      echo -e "  -h\tprint this help screen"
      echo -e "  -p\trequest credentials using the given awscli profile instead of the default one"
      echo -e "  -r\trequest credentials for the given IAM role ARN - can be combined with option '-p'"
      return 0
      ;;
    p)
      AWSCLI_OPTIONS=( "${AWSCLI_OPTIONS[@]}" "--profile ${OPTARG}" )
      case ${OPTARG} in
        profile1)
          AWSMFA_SERIAL="arn:aws:iam::111222333444:mfa/your.user"
          ;;
        profile2)
          AWSMFA_SERIAL="arn:aws:iam::555666777888:mfa/your.other-user"
          ;;
        profile3)
          unset AWSMFA_SERIAL
          ;;
      esac
      ;;
    r)
      AWSCLI_STS_COMMAND="assume-role"
      AWSCLI_STS_OPTIONS=( "${AWSCLI_STS_OPTIONS[@]}" "--role-arn ${OPTARG}" )
      AWSCLI_STS_OPTIONS=( "${AWSCLI_STS_OPTIONS[@]}" "--role-session-name ${USER}" )
      ;;
    \?)
      echo "Invalid option: -${OPTARG}" >&2
      return 1
      ;;
    :)
      unset OPTARG_DESCRIPTION

      case ${OPTARG} in
        p)
          OPTARG_DESCRIPTION="aws-cli profile"
          ;;
        r)
          OPTARG_DESCRIPTION="aws iam role-arn"
          ;;
      esac

      echo "Error: Option -${OPTARG} requires an argument: '${OPTARG_DESCRIPTION}'" >&2
      return 1
      ;;
  esac
done

if [ -n "${AWSMFA_SERIAL}" ]
then
  read -s -p "Enter MFA code: " token_code && echo
  AWSMFA_CREDENTIALS=( "--serial-number ${AWSMFA_SERIAL}" "--token-code ${token_code}" )
fi

reset-session-token
aws_session_information=$(aws ${AWSCLI_OPTIONS[@]} sts ${AWSCLI_STS_COMMAND} ${AWSCLI_STS_OPTIONS[@]} ${AWSMFA_CREDENTIALS[@]})

export AWS_ACCESS_KEY_ID=$(echo ${aws_session_information} | jq -r .Credentials.AccessKeyId)
export AWS_SECRET_ACCESS_KEY=$(echo ${aws_session_information} | jq -r .Credentials.SecretAccessKey)
export AWS_SESSION_TOKEN=$(echo ${aws_session_information} | jq -r .Credentials.SessionToken)
}