#!/bin/bash

DT_FORMAT="%Y-%m-%dT%H:%M:%S%z"
SESSION_EXPIRATION_TIMEOUT=3600  # 1 hour
SESSION_EXPIRATION_KEY="x_security_token_expires"


function urlencode() {
    local LC_ALL=C
    local c i n=${#1}
    for (( i=0; i<n; i++ )); do
        c="${1:i:1}"
        case "$c" in
            [-_.~A-Za-z0-9])
               printf '%s' "$c" ;;
            *) printf '%%%02X' "'$c" ;;
        esac
    done
    echo
}


function parse_aws_dt() {
  echo $(date -j -u -f ${DT_FORMAT} "${1::${#1}-3}${1:${#1}-2:${#1}}" '+%s')
}


function get_aws_session_expiration() {
  local now=$(date -u '+%s')
  local session_expiration_raw=$(aws configure get default.${SESSION_EXPIRATION_KEY})
  if ((${#session_expiration_raw} == 0)); then
    echo -1
  else
    local session_expiration=$(parse_aws_dt $session_expiration_raw)
    echo $((session_expiration-now))
  fi
}


function assume_aws_role() {
  local profile=${1}

  read -p "Enter MFA token: " mfa_token

  local source_profile=$(aws configure get ${profile}.source_profile)
  local role_arn=$(aws configure get ${profile}.role_arn)
  local mfa_serial=$(aws configure get ${profile}.mfa_serial)
  local region=$(aws configure get ${profile}.region)

  local creds_json=$(aws sts assume-role --output json --profile ${source_profile} --role-arn ${role_arn} --duration-seconds 43200 --role-session-name assumed-session --serial-number ${mfa_serial} --token-code ${mfa_token})

  aws configure set default.aws_access_key_id $(echo ${creds_json} | jq -r .Credentials.AccessKeyId)
  aws configure set default.aws_secret_access_key $(echo ${creds_json} | jq -r .Credentials.SecretAccessKey)
  aws configure set default.aws_session_token $(echo ${creds_json} | jq -r .Credentials.SessionToken)
  aws configure set default.${SESSION_EXPIRATION_KEY} $(echo ${creds_json} | jq -r .Credentials.Expiration)
  aws configure set default.region ${region}
}


function set_env_variables() {
  local aws_region_name=$(aws configure get default.region)
  local aws_access_key_id=$(urlencode $(aws configure get default.aws_access_key_id))
  local aws_secret_access_key=$(urlencode $(aws configure get default.aws_secret_access_key))
  local aws_session_token=$(urlencode $(aws configure get default.aws_session_token))

  export AIRFLOW_CONN_AWS_CONN="aws://${aws_access_key_id}:${aws_secret_access_key}@/?region_name=${aws_region_name}&aws_session_token=${aws_session_token}"
  export AIRFLOW_UID=$(id -u)
}


function renew_aws_session() {
  local profile=${1:-"gdb-whitelist-dev"}

  local session_expiration=$(get_aws_session_expiration)
  if ((${session_expiration} < ${SESSION_EXPIRATION_TIMEOUT})); then
    if ((${session_expiration} < 0)); then
      echo "AWS session expired for ${profile}"
    else
      echo "AWS session almost expired for ${profile}"
    fi
    echo "Renewing..."
    assume_aws_role ${profile}
  fi
}