assume

# Works in bash, zsh untested

function assume() {
        SOURCE_IDENTITY=""
        TAG=""

        RST="\033[0m"
        GRN="\033[32m"
        RED="\033[31m"


        PARGS=()

        while [[ $# -gt 0 ]]; do
          case $1 in
            -t|--tag)
                TAG="$2"
                shift # past argument
                shift # past value
                ;;
            -s|--source-identity)
                SOURCE_IDENTITY="$2"
                shift # past argument
                shift # past value
                ;;
            -*|--*)
                echo "Unknown option $1"
                exit 1
                ;;
            *)
                PARGS+=("$1") # save positional arg
                shift # past argument
                ;;
          esac
        done


        if [[ ${#PARGS[@]} -ne 2 ]]; then
                echo "Usage: $0 <role_arn> <session_name> [-s source_identity] [-t tag_key_value]"
                return
        fi

        TARGET_ARN="${PARGS[0]}"
        SESSION_NAME="${PARGS[1]}"

        cmd="aws sts assume-role --role-arn ${TARGET_ARN} --role-session-name ${SESSION_NAME}"

        if [[ -n "$TAG" ]]; then
                printf "${GRN}Setting transitive tag: ${TAG}\n" 1>&2
                TAG_KEY="$(echo $2|cut -d= -f1)"
                TAG_VALUE="$(echo $2|cut -d= -f2)"

                cmd="${cmd} --tags Key=${TAG_KEY},Value=${TAG_VALUE}"
        fi

        if [[ -n "$SOURCE_IDENTITY" ]]; then
                printf "${GRN}Source identity: ${SOURCE_IDENTITY}\n" 1>&2
                cmd="${cmd} --source-identity ${SOURCE_IDENTITY}"
        fi

        printf "${GRN}Session Name: ${SESSION_NAME}\n" 1>&2


        resp="$(
                resp=$(${cmd} 2>/dev/null)
                echo $resp | jq -r '.Credentials|"AWS_ACCESS_KEY_ID=\(.AccessKeyId) AWS_SECRET_ACCESS_KEY=\(.SecretAccessKey) AWS_SESSION_TOKEN=\(.SessionToken) AWS_SECURITY_TOKEN=\(.SessionToken)"'

        )"


        src_arn=$(aws sts get-caller-identity|jq -r '.Arn'|tr -d '\n')

        if ! echo "$resp" | grep -q AWS_ACCESS_KEY_ID; then
                current="${RED}Assume Role Failed (${TARGET_ARN})${RST}" 1>&2
        else
                export ${resp}
                current="${GRN}$(aws sts get-caller-identity 2>/dev/null|jq -r '.Arn'|tr -d '\n')${RST}"
        fi

        printf "${GRN}${src_arn}${RST} -> ${current}\n"
}