assume-role-okta

#!/bin/bash

# START USAGE DOCUMENTATION
# assume-role-okta is a command line tool to help assume roles through a bastion account.
# Store your bastion account credentials here ~/.aws/credentials
#
# Usage: assume-role-okta [account_name] [role] [aws-region]
# account_name          account id or alias
#                       aliases stored in ~/.aws/accounts as JSON {"alias": account_id}
#                       [default 'default']
# role                  the role to assume into the account
#                       [default 'read']
# aws_region            region to assume into default set in ~/.aws/config
#
# END USAGE DOCUMENTATION

# Colors
RED='\033[0;31m'
YELLOW='\033[0;33m'
GREEN='\033[0;32m'
NC='\033[0m' # No Color

assume-role-okta(){

  #######
  # FUNCTIONS
  #######
  sts_assume(){
    ROLE_SESSION_ARGS=(--role-arn "${role_arn}")
    ROLE_SESSION_ARGS+=(--profile "${source_profile_config}")
    ROLE_SESSION_ARGS+=(--duration-seconds "${ROLE_SESSION_TIMEOUT}")
    ROLE_SESSION_ARGS+=(--role-session-name "${user_clean}")
    ROLE_SESSION=$(aws sts assume-role "${ROLE_SESSION_ARGS[@]}" --output json || echo "fail")
  }

  oktasume(){
    full_output=($(okta-aws-cli -e -b -z | cut -d "=" -f2 ))
    
    export AWS_ACCESS_KEY_ID=${full_output[1]}
    export AWS_SECRET_ACCESS_KEY=${full_output[2]}
    export AWS_SESSION_TOKEN=${full_output[3]}
    export AWS_SECURITY_TOKEN=${full_output[3]}
  }

  echo_out() {
    if [ -z "$silent" ]; then
      # If this is outputting to an eval statement we need to output a "echo" as well
      if [ -z "$OUTPUT_TO_EVAL" ]; then
        echo "$@"
      else
        echo "echo \"$*\";"
      fi
    fi
  }

  set_eks(){
    eks_clusters=$(aws eks list-clusters --query 'clusters[*]' --output text)
    cluster_array=($(echo $eks_clusters | tr "\t" "\n"))

    if [ ${#cluster_array[@]} -eq 1 ]; then
      echo_out "󱃾 Switch to $eks_clusters"
      export KUBECONFIG="$HOME/.kube/$AWS_ACCOUNT_ID-$eks_clusters"
      aws eks update-kubeconfig --name $eks_clusters --alias $eks_clusters
      chmod 0600 $KUBECONFIG
    elif [ ${#cluster_array[@]} -gt 1 ]; then
      echo_out "More than one EKS Cluster found. Select:"
      eks_menu "${#cluster_array[@]}" "${cluster_array[@]}"
    else
      echo_out "No EKS Cluster found. ignoring ..."
    fi
  }

  eks_menu (){
    arrsize=$1
    # echo "Size of array: $arrsize"
    # echo_out "${@:2}"
    select option in "${@:2}"; do
      if [ "$REPLY" -gt "$arrsize" ];
      then
        echo_out "Exiting..."
        break;
      elif [ 1 -le "$REPLY" ] && [ "$REPLY" -le $((arrsize)) ];
      then
        # echo "You selected $option which is option $REPLY"
        echo_out "󰠳 Switch to $option"
        export KUBECONFIG="$HOME/.kube/$AWS_ACCOUNT_ID-$option"
        aws eks update-kubeconfig --name $option --alias $option
        chmod 0600 $KUBECONFIG
        break;
      else
        echo_out "Incorrect Input: Select a number 1-$arrsize"
      fi
    done
  }

  export_variables(){
    export AWS_ACCESS_KEY_ID=$(echo "$ROLE_SESSION" | jq -r .Credentials.AccessKeyId)
    export AWS_SECRET_ACCESS_KEY=$(echo "$ROLE_SESSION" | jq -r .Credentials.SecretAccessKey)
    export AWS_SESSION_TOKEN=$(echo "$ROLE_SESSION" | jq -r .Credentials.SessionToken)
    export AWS_SESSION_END=$(echo "$ROLE_SESSION" | jq -r .Credentials.Expiration)
    export AWS_SECURITY_TOKEN=$AWS_SESSION_TOKEN
    export AWS_ACCOUNT_ID="$account_id"
    export AWS_ACCOUNT_NAME="$account_name"
    export AWS_ACCOUNT_ROLE="$role"
    export GEO_ENV="$account_name" # For GeoEngineer https://github.com/coinbase/geoengineer
    export AWS_STS_ROLE_ARN="$role_arn"
    # echo_out "Success! IAM session envars are exported."
    echo "${GREEN}󱇲 switched role${NC}"
  }

  set_region(){
    AWS_CONFIG_REGION="$(aws configure get region --profile ${account_name})"
    if [ -z  "$AWS_CONFIG_REGION" ]; then
      AWS_CONFIG_REGION="$(aws configure get region --profile ${default_profile})"
    fi
    if [ -z "$aws_region_input" ] && [ -z "$AWS_REGION" ] && [ -z "$AWS_DEFAULT_REGION" ] && [ -z "$AWS_CONFIG_REGION" ] && [ -z "$OUTPUT_TO_EVAL" ]; then
      echo -n "Assume Into Region [us-east-1]: "
      read -r region
      region=${region:-"us-east-1"}
    elif [ ! -z "$aws_region_input" ]; then
      # if there is a $aws_region_input then set to $aws_region_input
      echo_out "Region from parameter $aws_region_input"
      region="$aws_region_input"
    elif [ ! -z "$AWS_CONFIG_REGION" ]; then
      echo_out "Region from config $AWS_CONFIG_REGION"
      region=$AWS_CONFIG_REGION
    elif [ ! -z "$AWS_REGION" ]; then
      # if there is a $AWS_REGION then set to $AWS_REGION
      region="$AWS_REGION"
    elif [ ! -z "$AWS_DEFAULT_REGION" ]; then
      # if there is a $AWS_DEFAULT_REGION then set to $AWS_DEFAULT_REGION
      region="$AWS_DEFAULT_REGION"
    fi
    if [ -z "$region" ]; then
      echo_out "region not defined"
      return
    fi
    export AWS_REGION="$region"
    export AWS_DEFAULT_REGION="$region"
  }


  #######
  # PRE-CONDITIONS
  #######

  # requires authentication tool
  if [ "ASSUME_TOOL" = "okta-aws-cli" ]; then
    if ! hash okta-aws-cli 2>/dev/null; then
      echo_out "assume-role requires 'okta-aws-cli' to be installed"
      return
    fi
  else
    if ! hash gimme-aws-creds 2>/dev/null; then
      echo_out "assume-role requires 'gimme-aws-creds' to be installed"
      return
    fi
  fi

  # requires jq
  if ! hash jq 2>/dev/null; then
    echo_out "assume-role requires 'jq' to be installed"
    return
  fi

  # requires aws
  if ! hash aws 2>/dev/null; then
    echo_out "assume-role requires 'aws' to be installed"
    return
  fi

  # exports
  export AWS_SESSION_START
  export AWS_SESSION_ACCESS_KEY_ID
  export AWS_SESSION_SECRET_ACCESS_KEY
  export AWS_SESSION_SESSION_TOKEN
  export AWS_SESSION_SECURITY_TOKEN
  export AWS_ACCESS_KEY_ID
  export AWS_SECRET_ACCESS_KEY
  export AWS_SESSION_TOKEN
  export AWS_SECURITY_TOKEN
  export AWS_REGION
  export AWS_DEFAULT_REGION
  export AWS_ACCOUNT_ID
  export AWS_ACCOUNT_NAME
  export AWS_ACCOUNT_ROLE
  export GEO_ENV
  export AWS_PROFILE_ASSUME_ROLE

    # INPUTS

  read -r -d '' HELP <<EOF
  OPTIONS:
  -s    Enable silent mode
  -r    Override Region
  -h    Display this help
EOF

  OPTIND=1                # Reset in case getopts has been used previously in the shell
  invalid_options=();     # Array for invalid options

  while getopts ":shr:" opt; do
    #echo_out "Actual opt: $opt"
    case $opt in
      s)
        silent='true'
        ;;
      r)
        aws_region_input=$OPTARG
        ;;
      h)
        echo "$HELP"
        exit 1
        ;;
      \?)
        invalid_options+=($OPTARG)
        ;;
     *)
        invalid_options+=($OPTARG)
        ;;
    esac
  done

  # HANDLE INVALID OPTIONS
  if [ ${#invalid_options[@]} -ne 0 ]; then
    echo "Invalid option(s):" >&2
    for i in "${invalid_options[@]}"; do
      echo $i >&2
    done
    echo "" >&2
    echo "$HELP" >&2
    exit 1
  fi

  # SET $1 TO FIRST MASS ARGUMENT, $2 TO SECOND MASS ARGUMENT ETC
  shift $((OPTIND - 1))

  account_name_input=$1

  CONF_FILE=${CONF_FILE:-~/.aws/config}
  SESSION_TIMEOUT=43200
  ROLE_SESSION_TIMEOUT=3600
  DEFAULT_ACCOUNT="default"
  DEFAULT_ROLE="read"
  user_clean=$USER

  if [ -z  "$AWS_PAGER" ]; then
    export AWS_PAGER="cat"
  fi

  # Force use of ~/.aws/credentials file which contains aws login account
  unset AWS_ACCESS_KEY_ID
  unset AWS_SECRET_ACCESS_KEY
  unset AWS_SESSION_TOKEN
  unset AWS_SECURITY_TOKEN

  if [ "$AWS_ACCOUNT_NAME" ]; then
    DEFAULT_ACCOUNT=${AWS_ACCOUNT_NAME}
  fi

  # set account_name
  if [ -z "$account_name_input" ] && [ -z "$OUTPUT_TO_EVAL" ]; then
    echo -n "Assume Into Account [$DEFAULT_ACCOUNT]:"
    read -r account_name
    # default
    account_name=${account_name:-"$DEFAULT_ACCOUNT"}
  else
    account_name="$account_name_input"
  fi

  if [ -f "$CONF_FILE" ]; then
    role_input=`cat $CONF_FILE | sed -n "/^\[profile $account_name\]/,/^\[.*\]/p" | grep "role_arn" | cut -d "=" -f2 | awk -F'/' '{ print $2 }'`

    if [ ! -z "$role_input" ] && [ "$role_input" != '' ]; then
      account_id=`cat $CONF_FILE | sed -n "/^\[profile $account_name\]/,/^\[.*\]/p" | grep "role_arn" | cut -d "=" -f2 | awk -F':' '{ print $5 }'`
      role_arn=`aws configure get role_arn --profile $account_name`
      source_profile_config=`aws configure get source_profile --profile $account_name || echo ''`
      if [ ! -z "$source_profile_config" ] && [ "$source_profile_config" != '' ]; then
        echo_out "Using source profile $source_profile_config"
        default_profile=$source_profile_config
      fi
    else
      account_id="$account_name_input"
      echo_out "Using Input as Account-ID instead of name"
      echo_out "Trying to find Account-ID in config-file"
      find_account=`cat $CONF_FILE | grep "$account_id" | wc -l`
      if [[ $find_account == 1 ]]; then
        role_arn=`cat $CONF_FILE | grep "$account_id" | awk '{print $3 }'`
        role_input=`cat $CONF_FILE | grep "$account_id" | awk '{print $3 }' | awk -F'/' '{ print $2 }'`
        account_name=`tac ~/.aws/config | grep -B 20 "$account_id" | grep "\[profile" | sed "s/\]$//" | awk '{print $2 }' | head -n1`
      else
        echo_out "Account-ID not found"
      fi
    fi
    # echo_out $role_input, $account_id
  fi

  # Validate Account ID
  if [[ ! $account_id =~ ^[0-9]{12}$ ]]; then
    echo_out "account_id \"$account_id\" is incorrectly formatted AWS account id"
    return
  fi

  # set role
  if [ -z "$role_input" ] && [ -z "$OUTPUT_TO_EVAL" ]; then
    echo -n "Assume Into Role [$DEFAULT_ROLE]: "
    read -r role
    role=${role:-"$DEFAULT_ROLE"}
  else
    role="$role_input"
  fi
  if [ -z "$role" ]; then
    echo_out "role not defined"
    return
  fi
  if [ -z "$role_arn" ]; then
    role_arn="arn:aws:iam::${account_id}:role/${role}"
  fi

  # set region
  set_region

  sts_assume

  if [ "$ROLE_SESSION" = "fail" ]; then
    echo_out "${YELLOW}Failed to export session envars.${NC}"
    if [ ! -z "$YUBIKEY_OKTA_KEY" ]; then
      OKTA_MFA_CODE=$(ykman oath accounts code "$YUBIKEY_OKTA_KEY" | awk '{print $NF}')
      echo_out "Using MFA-Key $YUBIKEY_MFA"
		export OKTA_MFA_CODE
    fi
    
    if [ "$ASSUME_TOOL" = "okta-aws-cli" ]; then
      echo_out "${YELLOW}Executing okta-aws-cli${NC}"
      oktasume
    else
      echo_out "${YELLOW}Executing gimme-aws-creds${NC}"
      if [ -z $GIMME_AWS_CREDS_CONF_PROFILE ]; then GIMME_AWS_CREDS_CONF_PROFILE="DEFAULT"; fi
      gimme-aws-creds -p $GIMME_AWS_CREDS_CONF_PROFILE
    fi

    sts_assume

    if [ "$ROLE_SESSION" = "fail" ]; then
      echo_out "${RED}Authentication did not work correctly try it manually${NC}"
    else
      export_variables
    fi

  else
    export_variables
  fi

  # Switch to EKS Cluster
  set_eks

  # OUTPUTS ALL THE EXPORTS for eval $(assume-role-okta [args])
  if [ "$OUTPUT_TO_EVAL" = "true" ]; then
    echo "export AWS_REGION=\"$AWS_REGION\";"
    echo "export AWS_DEFAULT_REGION=\"$AWS_DEFAULT_REGION\";"
    echo "export AWS_ACCESS_KEY_ID=\"$AWS_ACCESS_KEY_ID\";"
    echo "export AWS_SECRET_ACCESS_KEY=\"$AWS_SECRET_ACCESS_KEY\";"
    echo "export AWS_SESSION_TOKEN=\"$AWS_SESSION_TOKEN\";"
    echo "export AWS_ACCOUNT_ID=\"$AWS_ACCOUNT_ID\";"
    echo "export AWS_ACCOUNT_NAME=\"$AWS_ACCOUNT_NAME\";"
    echo "export AWS_ACCOUNT_ROLE=\"$AWS_ACCOUNT_ROLE\";"
    echo "export AWS_SESSION_ACCESS_KEY_ID=\"$AWS_SESSION_ACCESS_KEY_ID\";"
    echo "export AWS_SESSION_SECRET_ACCESS_KEY=\"$AWS_SESSION_SECRET_ACCESS_KEY\";"
    echo "export AWS_SESSION_SESSION_TOKEN=\"$AWS_SESSION_SESSION_TOKEN\";"
    echo "export AWS_SESSION_SECURITY_TOKEN=\"$AWS_SESSION_SESSION_TOKEN\";"
    echo "export AWS_SESSION_START=\"$AWS_SESSION_START\";"
    echo "export GEO_ENV=\"$GEO_ENV\";"
    echo "export AWS_PROFILE_ASSUME_ROLE=\"$AWS_PROFILE_ASSUME_ROLE\";"
    echo "export AWS_SECURITY_TOKEN=\"$AWS_SESSION_TOKEN\";"
    echo "export AWS_STS_ROLE_ARN=\"$role_arn\";"
    echo "export KUBECONFIG=\"$KUBECONFIG\";"
  fi

  # USED FOR TESTING AND DEBUGGING
  if [ "$DEBUG_ASSUME_ROLE" = "true" ]; then
    echo "AWS_CONFIG_REGION=\"$AWS_CONFIG_REGION\";"
    echo "SESSION_ARGS=\"${SESSION_ARGS[*]}\";"
    echo "SESSION=\"$SESSION\";"
    echo "ROLE_SESSION_ARGS=\"${ROLE_SESSION_ARGS[*]}\";"
    echo "ROLE_SESSION=\"$ROLE_SESSION\";"
    echo "SESSION_TIMEOUT=\"$SESSION_TIMEOUT\";"
    echo "ROLE_SESSION_TIMEOUT=\"$ROLE_SESSION_TIMEOUT\";"
    echo "AWS_PROFILE_ASSUME_ROLE=\"$AWS_PROFILE_ASSUME_ROLE\";"
    echo "AWS_STS_ROLE_ARN=\"$role_arn\";"
  fi

}



# from https://stackoverflow.com/questions/2683279/how-to-detect-if-a-script-is-being-sourced
if [[ "${BASH_SOURCE[0]}" != "${0}" ]]; then
  # If the script is being sourced, do nothing
  # Supports `source $(which assume_role)` in rc file for bash and zsh
  : # noop
elif [[ "init" == "${1}" ]]; then
  # TODO: This will be multi-shell support like rbenv, e.g. fish
  # Supports `eval "$(assume-role-okta init -)"` in rc file
  echo "Currently not supported"
else
  # The script is being called directly
  # Supports calling being called like eval $(assume-role-okta account role [token])
  set -eo pipefail
  OUTPUT_TO_EVAL="true"
  assume-role-okta "$@";
fi