Bash shell function to authenticate aws cli using mfa token

aws_mfa.sh

#!/bin/bash

function aws_mfa {
    code="$1"
    serial_number="<add mfa serial here>"

    # Check if token code is provided
    if [ -z "$code" ]; then
        echo -e "\e[31mToken code not provided. Please provide the token code as an argument.\e[0m"
        echo "Usage: aws_mfa <token-code> [additional arguments]"
        return 1
    fi

    # Run the AWS CLI command to get the session token
    sts_args=(--serial-number "$serial_number" --token-code "$code" "${@:2}")
    sts_command="aws sts get-session-token ${sts_args[@]}"
    echo "Executing AWS STS command:"
    echo "$sts_command"
    response=$(eval "$sts_command")

    # Check if the AWS CLI command succeeded
    if [ $? -ne 0 ]; then
        echo -e "\e[31mFailed to retrieve AWS session token. Please verify the token code and try again.\e[0m"
        return 1
    fi

    # Extract the necessary values from the response
    access_key=$(echo "$response" | jq -r '.Credentials.AccessKeyId')
    secret_key=$(echo "$response" | jq -r '.Credentials.SecretAccessKey')
    session_token=$(echo "$response" | jq -r '.Credentials.SessionToken')

    echo ""
    echo -e "\e[33mAWS_ACCESS_KEY_ID=$access_key"
    echo "AWS_SECRET_ACCESS_KEY=$secret_key"
    echo "AWS_SESSION_TOKEN=$session_token\e[0m"
    echo ""

    # Set the environment variables with the obtained credentials
    export AWS_ACCESS_KEY_ID="$access_key"
    export AWS_SECRET_ACCESS_KEY="$secret_key"
    export AWS_SESSION_TOKEN="$session_token"

    echo -e "\e[32mAWS credentials have been set.\e[0m"
}

# Example usage: aws_mfa <token-code> [additional arguments]