#!/usr/bin/env bash

unset LD_LIBRARY_PATH

function log {
  if [ -t 1 ]; then
    echo "$@" >&2
  fi
}

function prompt_password {
  if command -v kdialog >/dev/null; then
    kdialog --password Password
  else
    zenity --password --timeout=60
  fi
}

function prompt_mfa {
  if command -v kdialog >/dev/null; then
    kdialog --inputbox MFA
  else
    zenity --entry --text=MFA
  fi
}

# ignore existing credentials
unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN

output="eval"
duration=900
use_cache=1

eval set -- "$(getopt -o p:d:o: --long profile:,duration:,output: -- "$@")"
while true; do
  case "$1" in
    -p|--profile)
      profile=$2
      shift 2
      ;;
    -d|--duration)
      if [ -n "$2" ]; then
        duration=$2
        use_cache=0
      fi
      shift 2
      ;;
    -o|--output)
      if [ -n "$2" ]; then
        output=$2
      fi
      shift 2
      ;;
    --)
      shift
      break
      ;;
    *)
      echo "Internal error!"
      exit 1
  esac
done

cache_key="aws-$profile"
cache_gpg="$XDG_RUNTIME_DIR/$cache_key.gpg"

function get_from_keepassxc {
  keepassxc-cli show -q "$KEEPASS_FILE" "$KEEPASS_AWS_ENTRY" -a "$1" <<< "$2"
}

function get_role_arn {
  if ! aws configure get "profile.$profile.source_role_arn"; then
    log "source_role_arn not set for profile $profile"
    log "run aws configure set profile.$profile.source_role_arn <role arn>"
    return 1
  fi
}

function get_cached_long_term_credentials {
  if [ ! -f "$cache_gpg" ]; then
    log "No cached long-term credentials, requesting new"
    return 1
  fi

  if ! cached=$(gpg --batch -d --passphrase-file <(echo "$password") "$cache_gpg" 2>/dev/null); then
    log "Error getting cached long-term credentials"
    return 2
  fi

  expiration=$(date -d "$(jq -r '.Credentials.Expiration' <<< "$cached")" +%s)
  if [ "$expiration" -lt "$(date +%s)" ]; then
    log "Cached long-term credentials expired, requesting new"
    return 1
  fi

  echo "$cached"
}

function get_long_term_credentials {
  local duration
  local password
  duration="$1"
  if ! password=$(prompt_password 2> /dev/null); then
    log "Failed to get password"
    return 1
  fi

  credentials=$(get_cached_long_term_credentials)
  retcode=$?
  if [ $retcode -eq 2 ]; then
    exit 1
  elif [ $retcode -eq 1 ]; then
    if ! role_arn=$(get_role_arn); then
      return 1
    fi

    if ! AWS_ACCESS_KEY_ID=$(get_from_keepassxc UserName "$password"); then
      log "Failed to open vault"
      exit 1
    fi
    if ! AWS_SECRET_ACCESS_KEY=$(get_from_keepassxc Password "$password"); then
      log "Failed to open vault"
      exit 1
    fi

    export AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY
    export AWS_PROFILE=default

    mfa_serial_number=$(aws configure get mfa_serial)

    if ! mfa=$(prompt_mfa 2> /dev/null); then
      log "Failed to get MFA"
      return 1
    fi

    if credentials=$(aws sts assume-role --serial-number "$mfa_serial_number" --token-code "$mfa" --role-arn "$role_arn" --role-session-name "$(hostname)" --duration-seconds "$duration"); then
      gpg --batch -c --passphrase-file <(echo "$password") <<< "$credentials" > "$cache_gpg"
    else
      return 1
    fi
  fi

  echo "$credentials"
}

function get_cached_short_term_credentials {
  if [ "$use_cache" -eq 0 ]; then
    return 1
  fi

  keyctl pipe "%user:$cache_key" 2>/dev/null
}

function get_short_term_credentials {
  if ! credentials=$(get_cached_short_term_credentials); then
    if ! role_arn=$(get_role_arn); then
      return 1
    fi

    if ! long_term_credentials="$(get_long_term_credentials 43200)"; then
      return 1
    fi

    AWS_ACCESS_KEY_ID=$(jq -r .Credentials.AccessKeyId <<<"$long_term_credentials")
    AWS_SECRET_ACCESS_KEY=$(jq -r .Credentials.SecretAccessKey <<<"$long_term_credentials")
    AWS_SESSION_TOKEN=$(jq -r .Credentials.SessionToken <<<"$long_term_credentials")

    export AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN

    if credentials=$(aws sts assume-role --role-arn "$role_arn" --role-session-name "$(hostname)" --duration-seconds "$duration"); then
      keyctl padd user "$cache_key" @s <<<"$credentials" >/dev/null
      keyctl timeout "%user:$cache_key" "$duration" >/dev/null
    else
      return 1
    fi
  fi

  echo "$credentials"
}

while keyctl read '%user:lock' >/dev/null 2>&1; do
  sleep 1
done
keyctl add user lock lock @s >/dev/null 2>&1

function finish {
  keyctl revoke '%user:lock' 2>&1
}

trap finish EXIT

if [ "$duration" -gt 3600 ]; then
  if ! credentials=$(get_long_term_credentials "$duration"); then
    exit 1
  fi
else
  if ! credentials=$(get_short_term_credentials); then
    exit 1
  fi
fi

if [ "$output" = "credential_process" ]; then
  jq '.Credentials | .Version=1' <<< "$credentials"
else
  jq <<< "$credentials" -r '.Credentials | @sh "
export AWS_ACCESS_KEY_ID=\(.AccessKeyId)
export AWS_SECRET_ACCESS_KEY=\(.SecretAccessKey)
export AWS_SESSION_TOKEN=\(.SessionToken)"'
fi