#!/usr/bin/env bash
set -e
set -o pipefail

# FORMATTING ###################################################################

DEFAULT=`tput sgr0`
BOLD=`tput bold`
RED=`tput setaf 1`
GREEN=`tput setaf 2`
YELLOW=`tput setaf 3`
CYAN=`tput setaf 6`

# FUNCTIONS ####################################################################

# log - log to console with color output
function log {
  if [ $# -lt 2 ]; then
    local log_msg="$1"
  else
    local log_type="$1"
    local log_msg="$2"
  fi
  case "$log_type" in
    info) echo -e "\n${BOLD}${CYAN}${log_msg}${DEFAULT}" ;;
    error) echo -e "\n${BOLD}${RED}ERROR: ${log_msg}${DEFAULT}\n" ;;
    *) echo -e "\n${log_msg}" ;;
  esac
}

function usage {
  echo "Usage:"
  echo "  $(basname $0) <role> [cmd]   Run AWS command as role"
  echo "  $(basname $0) <role>         Generate env vars for assuming role"
  echo ""
  echo "Examples:"
  echo "  \$ $(basename $0) dev-deploy aws s3 ls"  # Run AWS command as dev-deploy
  echo "  \$ $(basename $0) operations-admin"      # Print out env vars for operations-admin
  echo ""
}

function lookup_account {
  ar_env=$1

  account=$(sed -n -e '/variable "aws_accounts"/,/}/ p' ${env_root}/variables.tf  | \
    sed -nE -e "s/[[:space:]]*${ar_env}[[:space:]]*=[[:space:]]*\"(.*)\"/\1/p"
  )

  if [[ -z $account ]]; then
      log error "Failed to lookup account for ${ar_env}" 1>&2
      exit 1
  fi

  log info $account
}

function aws_assume_role {
  ar_env=$1
  ar_role=$2

  if [[ $ar_env == "identity" ]]; then
    log info "Do not need to assume_role for the ops environment" 1>&2
    return
  fi

  account=$(lookup_account $ar_env)

  role="arn:aws:iam::${account}:role/${ar_role}"
  aws_tmp=$(mktemp -t aws-XXXX.json)

  aws sts assume-role --role-arn ${role} --role-session-name terraform > ${aws_tmp}

  aws_key=$(cat ${aws_tmp} | jq -r ".Credentials.AccessKeyId")
  aws_secret=$(cat ${aws_tmp} | jq -r ".Credentials.SecretAccessKey")
  aws_session_token=$(cat ${aws_tmp} | jq -r ".Credentials.SessionToken")
  aws_session_expiration=$(cat ${aws_tmp} | jq -r ".Credentials.Expiration")
}

function discover_aws_credentials {
  aws_region=${AWS_DEFAULT_REGION:-$(aws configure get region)}
  aws_key=${AWS_ACCESS_KEY_ID:-$(aws configure get aws_access_key_id)}
  aws_secret=${AWS_SECRET_ACCESS_KEY:-$(aws configure get aws_secret_access_key)}

  if [[ -z $aws_region || -z $aws_key || -z $aws_secret ]]; then
    log error "Could not get AWS credentials" 1>&2
    log info "Run 'aws configure' or set AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY/AWS_DEFAULT_REGION" 1>&2
    exit 1
  fi
}

function aws_env {
  if [[ -n $aws_session_token ]]; then
    export AWS_SESSION_TOKEN="$aws_session_token"
  fi
  export AWS_ACCESS_KEY_ID="$aws_key"
  export AWS_SECRET_ACCESS_KEY="$aws_secret"
  export AWS_DEFAULT_REGION="$aws_region"
  export TF_VAR_aws_region="$aws_region"
}

# Get ARGS
ENV=${1}
ROLE=${2:-deployer}
ARGS=${@:3}

if [[ $# < 1 ]]; then
    log error "Missing arguments"
    usage && exit 1
fi

unset AWS_DEFAULT_REGION \
      AWS_ACCESS_KEY_ID \
      AWS_SECRET_ACCESS_KEY \
      AWS_SESSION_TOKEN \
      AWS_SESSION_EXPIRATION

discover_aws_credentials
aws_assume_role $ENV $ROLE

if [[ $# > 0 ]]; then
  AWS_DEFAULT_REGION="$aws_region" \
  AWS_ACCESS_KEY_ID="$aws_key" \
  AWS_SECRET_ACCESS_KEY="$aws_secret" \
  AWS_SESSION_TOKEN="$aws_session_token" \
  AWS_SESSION_EXPIRATION="$aws_session_expiration" \
  "$@"
else
  echo export AWS_DEFAULT_REGION=\"$aws_region\"
  echo export AWS_ACCESS_KEY_ID=\"$aws_key\"
  echo export AWS_SECRET_ACCESS_KEY=\"$aws_secret\"
  echo export AWS_SESSION_TOKEN=\"$aws_session_token\"
  echo export AWS_SESSION_EXPIRATION=\"$aws_session_expiration\"
fi