Script to retrieve temp auth token from AWS STS

aws-session

#!/bin/bash

# Uses AWS CLI and user-provided MFA code to retrieve and store
# temp session credentials from AWS Security Token Service (STS),
# - for user sessions via get-session-token
# -  or role sessions via assume-role
#
# Stores 'token_expiration' in the profile to enable checks on session time left


# Details
#
# First run will prompt user to populate ARNs for MFA and a dev role:
#   aws --profile default configure set mfa_serial <arn-from-iam-console>
#   aws --profile default configure set <role-name> <arn-from-admin>
#
# Each session request will update a field named 'token_expiration' in the profile:
#
#   aws --profile $SESSION_PROFILE configure set token_expiration "$expiry"
#
# AWS CLI stores config values between two files: ~/.aws/config and ~/.aws/credentials.
#
# Long-term user credentials are distinguished from key-pair credentials for sessions
# by their prefix (AKIA vs ASIA):
#
#   AKIA - long-term key pair
#    - stored under [default] profile, with region and output format https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-quickstart.html
#   ASIA - temp key pair and session token, requested by authority of long-term access key pair
#    - from get-session-token https://awscli.amazonaws.com/v2/documentation/api/latest/reference/sts/get-session-token.html
#    - from assume-role https://awscli.amazonaws.com/v2/documentation/api/latest/reference/sts/assume-role.html
#
# This script assumes the following structure in the aws credentials and config files:
#   - all values stored under [$SOURCE_PROFILE] ('default', unless overridden with --source)
#     are considered read-only by this script; the user will need to set them ahead of time.
#   - this script creates and updates user and role sessions
#   - user and role sessions get their region from the default profile
#   - role sessions get their arn from the default profile by looking up the value of the provided role name
#   - configuration keys not listed here are ignored and left untouched
#
#                > credentials                   > config
#
#      long-term   [default]            _          [default]
#                  aws_access_key_id = AKIA..      region
#                  aws_secret_access_key           mfa_serial
#                                                  <role_name> = arn:aws:iam::<acct_id>:role/<role_name>
#
#   user session   [<user_name>]        _          [<user_name>]
#                  aws_access_key_id = ASIA..      region
#                  aws_secret_access_key           token_expiration
#                  aws_session_token
#
#   role session   [<role_name>]        _          [<role_name>]
#                  aws_access_key_id = ASIA..      region
#                  aws_secret_access_key           token_expiration
#                  aws_session_token               role_arn

DURATION_SEC=$(( 60 * 60 * 12 )) # default of twelve hours worth of seconds (43200)
SOURCE_PROFILE='default' # profile w/ long-term credentials to make requests from
SESSION_PROFILE='' # user-provided profile to save temp session credentials under
PROFILES=()
AWS_CLI=''
ARN_OF_ROLE=''
ARN_OF_MFA=''
MFA_CODE=''


#####

function out() { printf "%s\n" "$*" >&1; } # stdout for output to be consumed by programs
function msg() { printf "%b\n" "$*" >&2; } # stderr for messaging to be read by users; may use color if [ -t 2 ] is True

function usage() {
cat <<EOM

  request temp session token from aws security token service (sts),
  for a user (get-session-token), or role (assume-role).

  sessions will be requested using credentials in $SOURCE_PROFILE profile,
  use --source <name> to change

  usage:
    $(basename "$0") [args]

  args:
    --help
    --list
    --check|time <name>
    --user|role <name> [--source $SOURCE_PROFILE] [--duration $DURATION_SEC] <mfa-code>

  options
    -h --help             print this usage info
    -l --list             print names of known profiles, if any
    -c --check <name>     print status of session for <name> profile, if any
    -d --duration <sec>   (with -u or -r) set token life in seconds [900-129600] ($DURATION_SEC)
    -r --role <name>      assume role of arn stored in <name> in source profile
    -s --source <name>    request sessions using credentials in <name> profile ($SOURCE_PROFILE)
    -t --time <name>      print seconds remaining in session for <name> profile, if any
    -u --user <name>      request a session token without assuming a role

EOM
}

function exit_on_error() {
  local -r exit_code=$1
  shift
  [[ $exit_code ]] && ((exit_code != 0)) && {
    msg "$@"
    exit "$exit_code"
  }
}

function seconds_left() {
  local -r profile="$1"

  local -r timestamp=$($AWS_CLI --profile "$profile" configure get token_expiration)
  if [ -z "$timestamp" ]; then out "-1" && return 1; fi # no previous session detected (could be wrong profile)

  local -r seconds=$(( $(date '+%s' --date="$timestamp") - $(date '+%s') ))
  if [ "$seconds" -le 0 ]; then out "0" && return 0; fi # zero or negative seconds left; session expired

  out "$seconds" && return 0                            # positive seconds left; session still active
}

function session_status() {
  local -r profile="$1"
  local -r seconds=$(seconds_left "$profile")

  if [ "$seconds" -lt 0 ]; then out "no previous session detected for $profile" && return 1; fi
  if [ "$seconds" -eq 0 ]; then out "the previous $profile session has expired" && return 0; fi
  out "$profile has $((seconds / 60)) minutes left for the current session" && return 0
}

function clean_arn() {
  # avoid printing account numbers to logs and such
  # ARNs are colon delimited, with the most descriptive value in the last field
  local -r full_arn=$1
  out "$(awk -F: '{print $NF}' <<< "$full_arn")" && return 0
}

function read_profile_names() {
  if [ ${#PROFILES[@]} -gt 0 ]; then return 0; fi # assuming already read, since have entries
  IFS=$'\n' read -d '' -r -a PROFILES <<< "$($AWS_CLI configure list-profiles)"
}

function list_known_profiles() {
  read_profile_names
  if [ ${#PROFILES[@]} -gt 0 ]; then msg "known profiles include: $(awk -v OFS=", " '{$1=$1;print}' <<< "${PROFILES[*]}")"; fi
  return 0
}

function save_temp_creds() {
  if [ $# -ne 6 ]
  then
    msg "internal error: save_temp_creds() expected 6 arguments, got $#"
    msg ""
    return 1
  fi

  local -r profile=$1
  local -r resource=$2
  local -r id=$3
  local -r expiry=$4
  local -r key=$5
  local -r token=$6

  if [ "$profile" = "$SOURCE_PROFILE" ] || [ "$profile" = 'default' ]
  then
    msg "error: this could overwrite your long-term key pair"
    msg "       if this is something you really want to do,"
    msg "       please edit the credentials file directly."
    msg ""
    return 1
  fi

  if [ -z "$resource" ]
  then
    msg "error: no data received"
    msg ""
    return 1
  else
    msg "received $resource"
    msg ""
  fi

  msg "saving session data to --profile $profile.."
  local region; region="$(aws --profile $SOURCE_PROFILE configure get region)"; exit_on_error $? "could not read region from source profile '$SOURCE_PROFILE'"
  exit_on_error "$( \
    $AWS_CLI --profile "$profile" configure set region "$region" && \
    $AWS_CLI --profile "$profile" configure set aws_access_key_id "$id" && \
    $AWS_CLI --profile "$profile" configure set token_expiration "$expiry" && \
    $AWS_CLI --profile "$profile" configure set aws_secret_access_key "$key" && \
    $AWS_CLI --profile "$profile" configure set aws_session_token "$token"
  )" "failure calling aws configure set"

  return 0
}

function get_session() {
  msg "using credentials in '$SOURCE_PROFILE' profile to request temp session:"
  msg "  session: $SESSION_PROFILE"
  msg "  $(clean_arn "$ARN_OF_MFA"): $MFA_CODE"
  msg "  duration: $DURATION_SEC seconds"
  msg ""

  # https://awscli.amazonaws.com/v2/documentation/api/latest/reference/sts/get-session-token.html
  # https://awscli.amazonaws.com/v2/documentation/api/latest/reference/index.html#options
  local api_response
  api_response="$($AWS_CLI --profile $SOURCE_PROFILE \
    sts get-session-token \
      --duration $DURATION_SEC  \
      --serial-number "$ARN_OF_MFA" \
      --token-code "$MFA_CODE" \
      --output text)"
  exit_on_error $? "failure calling sts get-session-token"

  local resource=''
  local id=''
  local expiry=''
  local key=''
  local token=''

  # text output from api call gives the following space delimited structure (5 tokens):
  #   CREDENTIALS <aws_access_key_id> <token_expiration_timestamp> <aws_secret_access_key> <aws_session_token>
  #   ^ resource   ^ id                ^ expiry                     ^ key                   ^ token
  IFS=" " read -r resource id expiry key token <<< "$(awk '{ print $1, $2, $3, $4, $5 }' <<< "$api_response")"
  save_temp_creds "$SESSION_PROFILE" "$resource" "$id" "$expiry" "$key" "$token"
  exit_on_error $? "failure saving temp credentials"

  return 0
}

function assume_role() {
  msg "using credentials in '$SOURCE_PROFILE' profile to assume role:"
  msg "  $(clean_arn "$ARN_OF_ROLE")"
  msg "  $(clean_arn "$ARN_OF_MFA"): $MFA_CODE"
  msg "  duration: $DURATION_SEC seconds"
  msg ""

  # https://awscli.amazonaws.com/v2/documentation/api/latest/reference/sts/assume-role.html
  # https://awscli.amazonaws.com/v2/documentation/api/latest/reference/index.html#options
  local api_response
  api_response="$($AWS_CLI --profile $SOURCE_PROFILE \
    sts assume-role \
      --role-arn "$ARN_OF_ROLE" \
      --role-session-name "$USER-dev-session" \
      --duration-seconds $DURATION_SEC  \
      --serial-number "$ARN_OF_MFA" \
      --token-code "$MFA_CODE" \
      --output text)"
  exit_on_error $? "failure calling aws sts assume-role"

  local resource=''
  local id=''
  local expiry=''
  local key=''
  local token=''

  # text output from api call gives the following space delimited structure;
  # we ignore line one and take 5 tokens from line two:
  #   ASSUMEDROLEUSER <role_arn>        <session_name>
  #   CREDENTIALS <aws_access_key_id> <token_expiration_timestamp> <aws_secret_access_key> <aws_session_token>
  #   ^ resource   ^ id                ^ expiry                     ^ key                   ^ token
  IFS=" " read -r resource id expiry key token <<< "$(awk 'NR==2 {print $1, $2, $3, $4, $5}' <<< "$api_response")"
  save_temp_creds "$SESSION_PROFILE" "$resource" "$id" "$expiry" "$key" "$token"
  exit_on_error $? "failure saving temp credentials"

  return 0
}

#####

AWS_CLI=$(which aws)

if [ -z "$AWS_CLI" ]
then
  msg ""
  msg "error: cannot find AWS CLI on path"
  msg "  for documentation and installation help, see"
  msg "  docs.aws.amazon.com/cli/latest/userguide/"
  msg ""
  exit 1
fi

if [ $# -eq 0 ]; then usage && exit 2; fi

CMD=''
SHH=0
while [ $# -gt 0 ]
do
  opt="$1"
  case $opt in
    -h|--help)
      shift # past option
      usage
      exit 0
      ;;

    -l|--list)
      shift # past option
      CMD='list'
      SHH=1
      ;;

    -t|--time)
      SESSION_PROFILE="$2"
      shift # past option
      shift # past value
      CMD='time'
      SHH=1
      ;;

    -c|--check)
      SESSION_PROFILE="$2"
      shift # past option
      shift # past value
      CMD='check'
      SHH=1
      ;;

    -r|--role)
      SESSION_PROFILE="$2"
      shift # past option
      shift # past value
      CMD='role'
      ;;

    -u|--user)
      SESSION_PROFILE="$2"
      shift # past option
      shift # past value
      CMD='user'
      ;;

    -d|--duration)
      DURATION_SEC="$2"
      shift # past option
      shift # past value
      ;;

    -s|--source)
      SOURCE_PROFILE="$2"
      shift # past option
      shift # past value
      ;;

    *)    # unmatched option..
      if [ -z "${opt//[0-9]}" ]
      then  # if it's all digits, assume it's the required mfa code
        MFA_CODE="$opt"
        shift # past value
      else
        msg ""
        msg "error: unknown option '$opt'"
        msg ""
        exit 1
      fi
      ;;
  esac
done

ARN_OF_MFA=$($AWS_CLI --profile "$SOURCE_PROFILE" configure get mfa_serial)
if [ $? -gt 0 ]; then msg "failure calling aws configure get for field 'mfa_serial' in profile '$SOURCE_PROFILE'"; fi

if [ -z "$ARN_OF_MFA" ]
then
  msg ""
  msg "error: no mfa arn found for profile '$SOURCE_PROFILE'"
  msg "  find the amazon resource name (arn) for your mfa device at:"
  msg "  console.aws.amazon.com > user drop-down > my security credentials"
  msg "  then save it to your profile configuration:"
  msg ""
  msg "  aws --profile $SOURCE_PROFILE configure set mfa_serial <arn-from-iam-console>"
  msg ""
  exit 1
fi

if [ "$CMD" = 'time' ] || [ "$CMD" = 'check' ] || [ "$CMD" = 'role' ] || [ "$CMD" = 'user' ]
then
  if [ -z "$SESSION_PROFILE" ]
  then
    read_profile_names
    msg ""
    msg "error: missing name of session profile"
    msg "  you can:"
    if [ ${#PROFILES[@]} -gt 0 ]; then msg "  - use an existing name: (${PROFILES[*]})"; fi
    msg "  - provide a new name to have a profile created"
    msg ""
    exit 1
  fi
fi

if [ "$CMD" = 'role' ] || [ "$CMD" = 'user' ]
then
  if [ -z "$MFA_CODE" ]
  then
    msg ""
    msg "error: missing mfa token code"
    msg ""
    exit 1
  fi
fi

if [ "$CMD" = 'role' ]
then
  role_arn_field="$SESSION_PROFILE"
  ARN_OF_ROLE=$($AWS_CLI --profile "$SOURCE_PROFILE" configure get "$role_arn_field")
  if [ $? -gt 0 ]; then msg "failure calling aws configure get for field '$role_arn_field' in profile '$SOURCE_PROFILE'"; fi

  if [ -z "$ARN_OF_ROLE" ]
  then
    msg ""
    msg "error: no arn for '$role_arn_field' found in profile '$SOURCE_PROFILE'"
    msg "  ask your project admin for the amazon resource name (arn) for the"
    msg "  role you need to assume, and save it to your profile:"
    msg ""
    msg "  aws --profile $SOURCE_PROFILE configure set $role_arn_field <arn-from-admin>"
    msg ""
    exit 1
  fi
fi

if [ $SHH -eq 0 ];
then
  msg ""
  msg "$(basename "$0")"
  msg "using AWS CLI found at $AWS_CLI .."
  msg ""
fi

case $CMD in
  list)
    read_profile_names
    out "${PROFILES[*]}"
    ;;

  time)
    seconds="$(seconds_left "$SESSION_PROFILE")"
    code=$?
    if [ $code -gt 0 ]; then list_known_profiles; fi
    out "$seconds"
    exit $code
    ;;

  check)
    status="$(session_status "$SESSION_PROFILE")"
    code=$?
    if [ $code -gt 0 ]; then list_known_profiles; fi
    out "$status"
    exit $code
    ;;

  role)
    assume_role
    ;;

  user)
    get_session
    ;;

  *)    # should not get here..
    msg ""
    msg "internal error: unmatched command: '$CMD'"
    msg ""
    exit 1
    ;;
esac

if [ $SHH -eq 0 ];
then
  msg "done."
  msg ""
fi

exit 0

usage

$ aws --profile default configure list
      Name                    Value             Type    Location
      ----                    -----             ----    --------
   profile                  default           manual    --profile
access_key     ****************XXXX shared-credentials-file
secret_key     ****************XXXX shared-credentials-file
    region                 XX-XXX-X      config-file    ~/.aws/config

$ aws --profile default configure set mfa_serial <arn-from-iam-console>
$ aws --profile default configure set <role-name> <arn-from-admin>

$ ./aws-session --help

  request temp session token from aws security token service (sts),
  for a user (get-session-token), or role (assume-role).

  sessions will be requested using credentials in default profile,
  use --source <name> to change

  usage:
    aws-session [args]

  args:
    --help
    --list
    --check|time <name>
    --user|role <name> [--source default] [--duration 43200] <mfa-code>

  options
    -h --help             print this usage info
    -l --list             print names of known profiles, if any
    -c --check <name>     print status of session for <name> profile, if any
    -d --duration <sec>   (with -u or -r) set token life in seconds [900-129600] (43200)
    -r --role <name>      assume role of arn stored in <name> in source profile
    -s --source <name>    request sessions using credentials in <name> profile (default)
    -t --time <name>      print seconds remaining in session for <name> profile, if any
    -u --user <name>      request a session token without assuming a role

$ ./aws-session --role my-role 654321

aws-session
using AWS CLI found at /usr/local/bin/aws ..

using credentials in 'default' profile to assume role:
  role/my-role
  mfa/<user>: 654321
  duration: 43200 seconds

received CREDENTIALS

saving session data to --profile my-role..
done.

$ ./aws-session --check my-role
my-role has 717 minutes left for the current session

$ ./aws-session --time my-role
42973

$ ./aws-session --list
default my-role

@ellemenno you saved me a lot of hours of work!! This script is amazing! 💜