assume-role.sh

#!/bin/sh

# Source: https://stackoverflow.com/a/28776166/1960601
sourced=0
if [ -n "$ZSH_VERSION" ]; then 
  case $ZSH_EVAL_CONTEXT in *:file) sourced=1;; esac
elif [ -n "$KSH_VERSION" ]; then
  [ "$(cd -- "$(dirname -- "$0")" && pwd -P)/$(basename -- "$0")" != "$(cd -- "$(dirname -- "${.sh.file}")" && pwd -P)/$(basename -- "${.sh.file}")" ] && sourced=1
elif [ -n "$BASH_VERSION" ]; then
  (return 0 2>/dev/null) && sourced=1 
else # All other shells: examine $0 for known shell binary filenames.
  # Detects `sh` and `dash`; add additional shell filenames as needed.
  case ${0##*/} in sh|-sh|dash|-dash) sourced=1;; esac
fi

if [ $sourced -ne 1 ]; then
  echo "assume-role.sh script needs to be sourced"
  exit 1
fi

helpmsg() {
  echo "assume-role.sh [option]"
  echo " "
  echo "options:"
  echo "-h, --help                 show brief help"
  echo "-a, --assume <role_arn>    assume the role. E.g. arn:aws:iam::123456789012:role/the-role"
  echo "-i, --assume-interactively assume the role interactively (fzf required)"
  echo "-r, --reject               reject the role"
}

assume() {
  ROLE_ARN=$1
  if [ -z "$ROLE_ARN" ]; then
    echo No role specified. Exiting.
    return
  fi
  echo "assuming $ROLE_ARN"
  CREDENTIALS=$(aws sts assume-role --role-arn "$ROLE_ARN" \
     --role-session-name "session-$USER-assume-role" \
     --query Credentials)
  if [ -z "$CREDENTIALS" ]; then
    echo Unable to assume role. Exiting.
    return
  fi
  AWS_ACCESS_KEY_ID=$(echo "$CREDENTIALS" | jq -r '.AccessKeyId')
  AWS_SECRET_ACCESS_KEY=$(echo "$CREDENTIALS" | jq -r '.SecretAccessKey')
  AWS_SESSION_TOKEN=$(echo "$CREDENTIALS" | jq -r '.SessionToken')
  export AWS_ACCESS_KEY_ID
  export AWS_SECRET_ACCESS_KEY
  export AWS_SESSION_TOKEN
  aws sts get-caller-identity
}

assume_interactively() {
  ROLES=$(aws iam list-roles --query 'Roles[*].Arn|[][@]' --output text)
  if [ -z "$ROLES" ]; then
    echo No role found. Exiting.
    return
  fi
  ROLE_ARN=$(echo "$ROLES" | fzf)
  assume "$ROLE_ARN"
}

reject() {
  unset AWS_ACCESS_KEY_ID
  unset AWS_SECRET_ACCESS_KEY
  unset AWS_SESSION_TOKEN
  aws sts get-caller-identity
}

if [ $# -lt 1 ]; then
    helpmsg
fi
while [ $# -gt 0 ]; do
  case "$1" in
    -h|--help)
      helpmsg
      break
      ;;
    -a|--assume)
      shift
      if test $# -gt 0; then
        assume "$1"
      else
        echo "no role arn specified"
      fi
      break
      ;;
    -i|--assume-interactively)
      assume_interactively
      break
      ;;
    -r|--reject)
      reject
      break
      ;;
    *)
      helpmsg
      break
      ;;
  esac
done