Shell helper function for AWS STS AssumeRole

assume-role-helper.sh

#!/bin/zsh

# Check if jq is installed
function check_jq() {
  if ! command -v jq &> /dev/null; then
    echo "Error: jq is not installed. Please install jq to use this script."
    return 1
  fi
}

# Check if aws-cli is installed
function check_aws_cli() {
  if ! command -v aws &> /dev/null; then
    echo "Error: aws-cli is not installed. Please install aws-cli to use this script."
    return 1
  fi
}

function assume_role () {
  while [ "$#" -gt 0 ]; do
    case $1 in
      -p|--profile) profile="$2"; shift ;;
      --session-name) session_name="$2"; shift ;;
      --role-arn) role_arn="$2"; shift ;;
      -h|--help) assume_role_print_usage; return 0 ;;
      *) echo "Unknown parameter passed: $1"; assume_role_print_usage; return 1 ;;
    esac
    shift
  done

  check_jq
  check_aws_cli

  if [ -z "$session_name" ] || [ -z "$role_arn" ]; then
    assume_role_print_usage
    return 1
  fi

  # cmd="aws sts assume-role --profile \"$profile\" --role-arn \"$role_arn\" --role-session-name \"$session_name\""
  cmd="aws sts assume-role"
  rest_cmd="--role-arn \"$role_arn\" --role-session-name \"$session_name\""
  if [ -n "$profile" ]; then
    rest_cmd="--profile \"$profile\" $rest_cmd"
  fi

  resp=$(eval ${cmd} ${rest_cmd})

  echo "export AWS_ACCESS_KEY_ID=$(echo $resp | jq -r '.Credentials.AccessKeyId')"
  echo "export AWS_SECRET_ACCESS_KEY=$(echo $resp | jq -r '.Credentials.SecretAccessKey')"
  echo "export AWS_SESSION_TOKEN=$(echo $resp | jq -r '.Credentials.SessionToken')"
}

function assume_role_print_usage() {
  echo "Usage: assume_role --profile <profile> --session-name <session-name> --role-arn <role-arn>"
  echo "Arguments:"
  echo "  --profile: The name of the profile to use. (optional)"
  echo "  --session-name: An identifier for the assumed role session. (required)"
  echo "  --role-arn: The Amazon Resource Name (ARN) of the role to assume. (required)"
  echo "Example:"
  echo "  assume_role --profile default --session-name example-session --role-arn arn:aws:iam::123456789012:role/example-role"
}

usage (example)

$ ls 
assume-role-helper.sh

$ mkdir -p ~/.zsh_functions

$ mv ./assume-role-helper.sh ~/.zsh_functions

$ cat <<'EOF' >> ~/.zshrc
for fn in ~/.zsh_functions/*.sh; do
  [ -f "$fn" ] && source "$fn"
done
EOF