#!/bin/bash

set -euo pipefail

: ${AWS_CREDENTIALS_FILE="$HOME/.aws/credentials"}
: ${STS_SESSION_DURATION=36000}

sts_keychain_get_session_token() {
  echo generating temporary credentials via sts >&2
  aws-keychain exec $1 \
    aws sts --output text get-session-token --duration-seconds ${STS_SESSION_DURATION}
  echo credentials are valid for ${STS_SESSION_DURATION}s >&2
}

sts_keychain_format_credentials() {
  local id="$1"
  local secret="$2"
  local token="$3"
  cat <<END
[default]
aws_access_key_id=$id
aws_secret_access_key=$secret
aws_session_token=$token
END
}

sts_keychain_format_env() {
  local id="$1"
  local secret="$2"
  local token="$3"
  cat <<END
export AWS_ACCESS_KEY_ID="$id"
export AWS_SECRET_ACCESS_KEY="$secret"
export AWS_SESSION_TOKEN="$token"
END
}

sts_keychain_exec() {
  local name="$2"
  shift 2
  eval $($0 env "$name"); exec "$@"
}

sts_keychain_cat() {
  local name="$2"
  local output=$(sts_keychain_get_session_token $name)
  local id=$(awk '{print $2}' <<< $output)
  local secret=$(awk '{print $4}' <<< $output)
  local token=$(awk '{print $5}' <<< $output)
  sts_keychain_format_credentials "$id" "$secret" "$token"
}

sts_keychain_env() {
  local name="$2"
  local output=$(sts_keychain_get_session_token $name)
  local id=$(awk '{print $2}' <<< $output)
  local secret=$(awk '{print $4}' <<< $output)
  local token=$(awk '{print $5}' <<< $output)
  sts_keychain_format_env "$id" "$secret" "$token"
}

sts_keychain_use() {
  local name="$2"
  sts_keychain_cat "$@" > $AWS_CREDENTIALS_FILE
}

case "${1:-}" in
  cat) sts_keychain_cat "$@"; exit 0 ;;
  env) sts_keychain_env "$@"; exit 0 ;;
  exec) sts_keychain_exec "$@"; exit 0 ;;
  use) sts_keychain_use "$@"; exit 0 ;;
esac

exec aws-keychain "$@"