#!/bin/sh

#GRANTED_FLAG - what assumego told the shell to do
#GRANTED_n - the data from assumego

# pass an environment variable to the Go binary if the Granted alias hasn't been configured
_this_type=$(type -- "${0##*/}" 2>&1)
# If the type command failed, this is probably a sourced
# file because $0 will be the login shell, e.g., -bash.
# In that case the output will typically contain 'not found'.
# In the case of zsh, the output will contain the word 'alias'.
# shellcheck disable=SC3028
if [ "${_this_type#*not found}" != "$_this_type" ] ||
    [ "${_this_type#*alias}" != "$_this_type" ] ||
    [ "${BASH_SOURCE:-$0}" != "${0}" ]; then
  GRANTED_RETURN_STATUS="true"
  export GRANTED_ALIAS_CONFIGURED="true"
fi

# in some cases such as github actions, the SHELL env var is not available to the program
# so here we set it explicitly
GRANTED_OUTPUT=$(assumego "$@")

GRANTED_STATUS=$?
# shellcheck disable=SC2162
IFS=' ' read GRANTED_FLAG GRANTED_1 GRANTED_2 GRANTED_3 GRANTED_4 GRANTED_5 GRANTED_6 GRANTED_7 GRANTED_8 GRANTED_9 GRANTED_10 GRANTED_11 GRANTED_12<< EOF
${GRANTED_OUTPUT}
EOF

# # unset the exported GRANTED_ALIAS_CONFIGURED flag
unset GRANTED_ALIAS_CONFIGURED

# remove carriage return
GRANTED_FLAG=$(printf '%s\n' "$GRANTED_FLAG" | tr -d '\r')

if [ "$GRANTED_FLAG" = "GrantedDesume" ]; then
  unset AWS_ACCESS_KEY_ID
  unset AWS_SECRET_ACCESS_KEY
  unset AWS_SESSION_TOKEN
  unset AWS_PROFILE
  unset AWS_REGION
  unset AWS_DEFAULT_REGION
  unset AWS_SESSION_EXPIRATION
  unset AWS_CREDENTIAL_EXPIRATION
  unset GRANTED_SSO
  unset GRANTED_SSO_START_URL
  unset GRANTED_SSO_ROLE_NAME
  unset GRANTED_SSO_REGION
  unset GRANTED_SSO_ACCOUNT_ID
fi


if [ "$GRANTED_FLAG" = "GrantedAssume" ]; then
  # unset any previous vars
  unset AWS_ACCESS_KEY_ID
  unset AWS_SECRET_ACCESS_KEY
  unset AWS_SESSION_TOKEN
  unset AWS_PROFILE
  unset AWS_REGION
  unset AWS_DEFAULT_REGION
  unset AWS_SESSION_EXPIRATION
  unset AWS_CREDENTIAL_EXPIRATION
  unset GRANTED_SSO
  unset GRANTED_SSO_START_URL
  unset GRANTED_SSO_ROLE_NAME
  unset GRANTED_SSO_REGION
  unset GRANTED_SSO_ACCOUNT_ID

  # shellcheck disable=SC2124
  export GRANTED_COMMAND="$@"

  if [ ! "${GRANTED_1}" = "None" ]; then
    export AWS_ACCESS_KEY_ID="${GRANTED_1}"
  fi
  if [ ! "${GRANTED_2}" = "None" ]; then
    export AWS_SECRET_ACCESS_KEY="${GRANTED_2}"
  fi
  if [ ! "${GRANTED_3}" = "None" ]; then
    export AWS_SESSION_TOKEN="${GRANTED_3}"
  fi
  if [ ! "${GRANTED_4}" = "None" ]; then
    export AWS_PROFILE="${GRANTED_4}"
  fi
  if [ ! "${GRANTED_5}" = "None" ]; then
    export AWS_REGION="${GRANTED_5}"
    export AWS_DEFAULT_REGION="${GRANTED_5}"
  fi
  # The following are both variations of the same thing supported by the AWS CLI
  # AWS_SESSION_EXPIRATION
  # AWS_CREDENTIAL_EXPIRATION
  if [ ! "${GRANTED_6}" = "None" ]; then
    export AWS_SESSION_EXPIRATION="${GRANTED_6}"
    export AWS_CREDENTIAL_EXPIRATION="${GRANTED_6}"
  fi
  if [ ! "${GRANTED_7}" = "None" ]; then
    export GRANTED_SSO="${GRANTED_7}"
  fi
  if [ ! "${GRANTED_8}" = "None" ]; then
    export GRANTED_SSO_START_URL="${GRANTED_8}"
  fi
  if [ ! "${GRANTED_9}" = "None" ]; then
    export GRANTED_SSO_ROLE_NAME="${GRANTED_9}"
  fi
  if [ ! "${GRANTED_10}" = "None" ]; then
    export GRANTED_SSO_REGION="${GRANTED_10}"
  fi
  if [ ! "${GRANTED_11}" = "None" ]; then
    export GRANTED_SSO_ACCOUNT_ID="${GRANTED_11}"
  fi
fi

# Mark: Automatically re-assume when credentials expire.
_is_assume_expired() {
  [ -z "${AWS_PROFILE}" ] && return 1

  # Note: this must remain compatible with both BSD and GNU date.
  # TODO: This should probably run a few minutes (configurable) before it expires.
  # shellcheck disable=SC2034,SC3043 # Only used by zsh
  local curent_time expiry
  current_time="$(date -Iseconds)"
  expiry="$AWS_SESSION_EXPIRATION"
  # shellcheck disable=SC3010,SC3012,SC3054 # Only used by zsh
  [[ "${current_time}" > "${expiry}" ]]
}

granted_auto_reassume() {
  # Nothing to do, we can't reassume a profile that we don't know.
  [ -z "${AWS_PROFILE}" ] && return 0

  _is_assume_expired || return 0

  [ "${GRANTED_QUIET:-}" = "true" ] ||
    printf 'granted session expired; reassuming %s\n.' "${AWS_PROFILE}" >&2
  assume "${AWS_PROFILE}"
}

if [ -n "${ZSH_NAME:-}" ] && [ "${GRANTED_ENABLE_AUTO_REASSUME:-}" = "true" ]; then
  # shellcheck disable=SC2154,SC3054
  pfuncs=$(print -l -- "${preexec_functions[*]}")
  if [ "${pfuncs#*granted_auto_reassume}" = "$pfuncs" ]
  then
    autoload -Uz add-zsh-hook
    add-zsh-hook preexec granted_auto_reassume
  fi
fi

# Execute an additional program when GRANTED_FLAG is GrantedExec
if [ "$GRANTED_FLAG" = "GrantedExec" ]; then
  # Set GRANTED_12 with a command to execute, for example, "bash -c 'some_command'"
  # Make sure to properly escape quotes if needed and pass arguments.

  # Set and export the AWS variables only for the duration of the 'sh -c' command
  AWS_ACCESS_KEY_ID="${GRANTED_1:-}" \
  AWS_SECRET_ACCESS_KEY="${GRANTED_2:-}" \
  AWS_SESSION_TOKEN="${GRANTED_3:-}" \
  AWS_REGION="${GRANTED_5:-}" \
  AWS_DEFAULT_REGION="${GRANTED_5:-}" \
  sh -c "$GRANTED_12"

  # The variables will not affect the parent shell environment after the 'sh -c' command completes
fi
# The GrantedOutput flag should be followed by a newline, then the output.
# This way, the shell script can omit the first line containing the flag and return the unaltered output to the stdout
# This is great as it works well with the -exec flag
if [ "$GRANTED_FLAG" = "GrantedOutput" ]; then
  printf '%s\n' "${GRANTED_OUTPUT}" | sed -n '1!p'
fi

if [ "$GRANTED_RETURN_STATUS" = "true" ]; then
  # We definitely want return here and not exit because we are operating within
  # a sourced file. Calling exit would exit the login shell.
  return $GRANTED_STATUS
fi
