aws_session

#! /usr/bin/zsh

# Simplify authorising with MFA and assuming a role using the AWS CLI.
#
# Required parameters:
#   --token | -t: MFA token
#
# Optional parameters:
#   --echo-vars | -e: If supplied, the script will echo the values of the AWS env variables after they have been set.
#
# Author: Jon Karlsen (jon.karlsen@azets.com), 08 July 2022 


#######################################
#   Print the date, time, and a supplied
#   error message to STDOUT, then exit with
#   status 1.
#   Globals:
#     None
#   Arguments:
#     None
#   Outputs:
#     None
#######################################
function err {
  echo ""
  echo "[$(date +'%Y-%m-%dT%H:%M:%S%z')]: $*" >&2
  return 1
}


#######################################
#   Unset AWS env variables to start auth 
#   with a clean slate. 
#   Globals:
#       None
#   Arguments:
#       None
#   Outputs:
#       None
#######################################
function prep {
    echo "Cleaning up env variables..."

    unset AWS_ACCESS_KEY_ID
    unset AWS_SECRET_ACCESS_KEY
    unset AWS_SESSION_TOKEN
    
    echo "...done."
}


#######################################
#   Read credentials from config file.
#   Globals:
#       None
#   Arguments:
#       None
#   Outputs:
#       None
#######################################
function read_credentials {
    readonly __cred_file=${1:?"The path to the credentials file is required"}

    echo ""
    echo "Attempting to read credentials from $__cred_file..."

    while IFS=' = ' read key value; do
        if [[ $key = \[*] ]]; then
            section=$key
        elif [[ $value ]] && [[ $section = '[default]' ]]; then
            if [[ $key = 'aws_access_key_id' ]];       then
                aws_access_key_id=$value
            elif [[ $key = 'aws_secret_access_key' ]]; then
                aws_secret_access_key=$value
            elif [[ $key = 'mfa_serial' ]];            then
                mfa_serial=$value
            fi
        fi
    done < $__cred_file

    echo "...done."
}


#######################################
#   Read credentials from config file.
#   Globals:
#       aws_access_key_id
#       aws_secret_access_key
#       mfa_serial
#       aws_role_arn
#       aws_role_session_name
#   Arguments:
#       None
#   Outputs:
#       None
#######################################
function verify_credentials {
    if   [ ${#aws_access_key_id}     = 0 ]; then err "Failed to read a value for the required 'aws_access_key_id' configuration option"
    elif [ ${#aws_secret_access_key} = 0 ]; then err "Failed to read a value for the required 'aws_secret_access_key' configuraton option"
    elif [ ${#mfa_serial}            = 0 ]; then err "Failed to read a value for the required 'mfa_serial' configuration option"
    fi
}


    token=""
echo_vars=false
conf_file="~/.aws/credentials"


for i in "$@"; do
    case $i in
        -t=* | --token=* )
            token="${i#*=}"
            shift
            ;;
        -e=* | --echo-vars=* )
            if [ "${i#*=}" = "true" ]; then
                echo_vars=true
            fi
            shift
            ;;
        -* |--* )
            err "Unknown option $i"
            ;;
    esac
done


if [ ${#token} = 0 ]; then 
    err "Failed to parse a value for the required 'token' option"
else
    echo ""

    prep
    read_credentials ~/.aws/credentials
    verify_credentials
    
    echo ""
    echo "Attempting to get a session token..."

    aws sts get-session-token \
        --serial-number "$mfa_serial" \
        --token-code "$token" \
        --output json \
        --duration-seconds 14400 \
        > /tmp/session_token.json

    echo "...done."

    echo ""
    echo "Attempting to read the session token response..."

        aws_access_key_id=$( cat "/tmp/session_token.json" | jq -r '.Credentials.AccessKeyId' )
    aws_secret_access_key=$( cat "/tmp/session_token.json" | jq -r '.Credentials.SecretAccessKey' )
        aws_session_token=$( cat "/tmp/session_token.json" | jq -r '.Credentials.SessionToken' )

    if   [ ${#aws_access_key_id}     = 0 ]; then err "Failed to read a value for the required 'AccessKeyId' field"
    elif [ ${#aws_secret_access_key} = 0 ]; then err "Failed to read a value for the required 'SecretAccessKey' field"
    elif [ ${#aws_session_token}     = 0 ]; then err "Failed to read a value for the required 'SessionToken' field"
    else 
          tmp_file=$(mktemp)
        in_section=""

        while read -r line;
        do
            if [[ "$line" =~ (\[[a-z ]+\]) ]]; then
                if [[ "${match[1]}" == "[mfa]" ]]; 
                    then in_section="1"
                    else in_section=""
                fi
            fi

            if [[ "$in_section" == "1" ]];
            then 
                if [[ "$line" =~ ^aws_access_key_id ]];     then line="aws_access_key_id    = $aws_access_key_id"
                fi

                if [[ "$line" =~ ^aws_secret_access_key ]]; then line="aws_secret_access_key = $aws_secret_access_key"
                fi

                if [[ "$line" =~ ^aws_session_token ]];     then line="aws_session_token    = $aws_session_token"
                fi
            fi

            echo "$line" >> "$tmp_file"

        done < ~/.aws/credentials

        mv "$tmp_file" ~/.aws/credentials

        if [ "$echo_vars" = true ]; then
            echo "AWS_ACCESS_KEY_ID: "     $aws_access_key_id
            echo "AWS_SECRET_ACCESS_KEY: " $aws_secret_access_key
            echo "AWS_SESSION_TOKEN: "     $aws_session_token
            echo ""
        fi

            aws_access_key_id=""
        aws_secret_access_key=""
            aws_session_token=""

        rm /tmp/session_token.json
    fi

    echo "...done."
    echo ""

    return 0
fi