aws-assume

#!/usr/bin/env bash
# Uses AWS STS to assume role, setting environment variables
# To run: . aws-assume

pkg=${0##*/}

## Functions
log() {
  echo "$(date +"[%F %X,000]") $(hostname) $1"
}

die() {
  log "[FAIL] $1" >&2 && return 1
}

role_unset() {
  # Unsets environment variables #
  unset AWS_ACCESS_KEY_ID
  unset AWS_SECRET_ACCESS_KEY
  unset AWS_SESSION_TOKEN
  unset AWS_SESSION_EXPIRY
  unset AWS_ROLE_NAME
}

role_status() {
  # Checks the role status and displays the expiration time remaining #
  TIMELEFT=0
  if [[ ! -z "${AWS_SESSION_EXPIRY}" ]]; then
    EXPEPOCH=$(date --date "${AWS_SESSION_EXPIRY}" +%s)
    NOWEPOCH=$(date +%s)
    TIMELEFT=$((${EXPEPOCH} - ${NOWEPOCH}))

    if [[ "$TIMELEFT" -gt 0 ]]; then
      log "Assumed IAM Role: ${AWS_ROLE_NAME} (expires in "${TIMELEFT}"s)"
    else
      role_unset
      log "Role of ${ROLE_ARN} expired."
    fi
  fi
}

usage() {
  # Show usage #
  echo "
NAME
  ${pkg} - AWS STS role assumption

DESCRIPTION
  ${pkg} will, based on parameters passed, assume a role and provide AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY
  and an AWS_SESSION_TOKEN which are needed for SDK, CLI, and other such access requirements for tools. You
  can also pass your MFA token if enabled on your profile.

SYNOPSIS
  ${pkg} [--profile=<name>] [--account=<id>] [--role=<name>] [--user=<username>] [--mfa=<token>]

OPTIONS
  --profile=<name>
    Name of AWS profile from your ~/.aws/config file to use.

  --account=<id>
    AWS Account ID in which operation will be performed.

  --role=<name>
    Name of AWS role to assume.

  --user=<username>
    Name of AWS IAM User which has the MFA device.

  --mfa=<token>
    MFA token.
"
}

assume() {
  # Assume the specified role and get the STS credentials to be set within the environment #
  if [ -z "${AWS_PROFILE_NAME}" ] || [ -z "${AWS_ROLE_NAME}" ]; then
    die "Missing parameters. Set both '--profile=<name>' and '--role=<name>', or environment variables 'AWS_PROFILE_NAME' and 'AWS_ROLE_NAME'." ; return 1
  fi
  ROLE_ARN=$(aws --profile=${AWS_PROFILE_NAME} --output=json iam get-role --role-name ${AWS_ROLE_NAME} 2>/dev/null | jq -r '.Role | .Arn')
  if [ $? -ne 0 ] || [ -z "${ROLE_ARN}" ]; then
    die "Error getting role ARN." ; return 1
  fi
  if [ -z "${AWS_ACCOUNT_ID}" ] || [ -z "${USER}" ]; then
    die "Missing parameters. Set both '--account=<id>' and '--user=<username>', or environment variables 'AWS_ACCOUNT_ID' and 'USER'." ; return 1
  fi
  MFA_SERIAL_NUMBER="arn:aws:iam::${AWS_ACCOUNT_ID}:mfa/${USER}"
##  AWS_SESSION=$(aws --profile=${AWS_PROFILE_NAME} sts assume-role --output=json --role-arn "${ROLE_ARN}" --role-session-name ${AWS_PROFILE_NAME}-${AWS_ROLE_NAME} --serial-number "${MFA_SERIAL_NUMBER}" --token-code "${MFA_TOKEN}")
  AWS_SESSION=$(aws --profile=${AWS_PROFILE_NAME} sts assume-role --output=json --role-arn "${ROLE_ARN}" --role-session-name ${AWS_PROFILE_NAME}-${AWS_ROLE_NAME} 2>/dev/null)
  if [ $? -eq 0 ] && [ ! -z "${AWS_SESSION}" ]; then
    export AWS_SECRET_ACCESS_KEY=$(echo "${AWS_SESSION}" | jq -r '.Credentials|.SecretAccessKey')
    export AWS_ACCESS_KEY_ID=$(echo "${AWS_SESSION}" | jq -r '.Credentials|.AccessKeyId')
    export AWS_SESSION_TOKEN=$(echo "${AWS_SESSION}" | jq -r '.Credentials|.SessionToken')
    export AWS_SESSION_EXPIRY=$(echo "${AWS_SESSION}" | jq -r '.Credentials|.Expiration')
    role_status
  else
    die "Error assuming role to get credentials." ; return 1
  fi
}

main () {
  # Main function to handle processing #
  # Process command line
  for arg in "${@}"; do
    case "${arg}" in
      --profile)
        if [ -n "${2}" ]; then
          AWS_PROFILE_NAME=${2}
          shift
        else
          die "Option ${arg} requires a value." ; return 1
        fi
      ;;
      --profile=*)
        AWS_PROFILE_NAME=${arg#*=} # Delete everything up to "=" and assign the remainder.
      ;;
      --account)
        if [ -n "${2}" ]; then
          AWS_ACCOUNT_ID=${2}
          shift
        else
          die "Option ${arg} requires a value." ; return 1
        fi
      ;;
      --account=*)
        AWS_ACCOUNT_ID=${arg#*=} # Delete everything up to "=" and assign the remainder.
      ;;
      --role)
        if [ -n "${2}" ]; then
          AWS_ROLE_NAME=${2}
          shift
        else
          die "Option ${arg} requires a value." ; return 1
        fi
      ;;
      --role=*)
        AWS_ROLE_NAME=${arg#*=} # Delete everything up to "=" and assign the remainder.
      ;;
      --mfa)
        if [ -n "${2}" ]; then
          MFA_TOKEN=${2}
          shift
        else
          die "Option ${arg} requires a value." ; return 1
        fi
      ;;
      --mfa=*)
        MFA_TOKEN=${arg#*=}
      ;;
      -h | --help)
        usage
        return 0
      ;;
      *)
        log "Invalid option '${arg}'"
        usage
        return 1
      ;;
    esac
  done

  assume
}

## Main
main $@