#!/bin/bash

declare -f kstore > /dev/null
if [ $? -ne 0 ]; then
	echo "kotp depends on kstore" >&2
	return 1
fi

which oathtool > /dev/null 2>&1
if [ $? -ne 0 ]; then
	echo "oathtool is required" >&2
	return 1
fi

_KOTP_PREFIX="${_KOTP_PREFIX:-otp}"
_KOTP_PROP="${_KOTP_PROP:-secret}"

function kotp {
	case "$1" in
		set) shift; _kotp-set "$@" ;;
		gen|get) shift; _kotp-gen "$@" ;;
		spread) shift; _kotp-spread "$@" ;;
		list) shift; _kotp-list "$@" ;;
		del|delete|rm) shift; _kotp-del "$@" ;;
		*)
			__func_head "set NAME [SECRET|-|file]    set/update otp key"
			__func_help "gen NAME [TIME]              generate otp at TIME; defaults to now"
			__func_help "spread NAME [COUNT] [TIME]   generate surrounding otp codes; COUNT defaults to 1"
			__func_help "list [TERM]                  list stored otp keys"
			__func_help "del NAME                     delete otp key"
			return 1
			;;
	esac
	return $?
}

function _kotp-key {
	echo "${_KOTP_PREFIX}:$1"
}

function _kotp-clean-secret {
	local s

	if [ $# -gt 0 ]; then
		s="$1"
	else
		s="$(cat)"
	fi

	case "$s" in
		otpauth://*)
			s="$(printf '%s' "$s" \
				| sed -n 's/.*[?&]secret=\([^&]*\).*/\1/p')"
			;;
	esac

	printf '%s' "$s" \
		| tr -d '[:space:]-' \
		| tr '[:lower:]' '[:upper:]'
}

function _kotp-read-value {
	if [ "$1" = "-" ]; then
		cat
	elif [ -n "$1" ] && [ -f "$1" ]; then
		cat "$1"
	else
		printf '%s' "$1"
	fi
}

function _kotp-set {
	if [ -z "$1" ]; then
		__func_head "NAME [SECRET|-|file]"
		return 1
	fi

	local name src secret key
	name="$1"
	src="$2"

	if [ -z "$src" ]; then
		read -rsp "OTP secret: " src
		echo
	fi

	secret="$(_kotp-read-value "$src" | _kotp-clean-secret)"
	if [ -z "$secret" ]; then
		echo "empty otp secret" >&2
		return 1
	fi

	key="$(_kotp-key "$name")"

	# Validate before storing. Most TOTP secrets are base32.
	oathtool --totp --base32 "$secret" > /dev/null 2>&1
	if [ $? -ne 0 ]; then
		echo "invalid otp secret, or oathtool cannot parse it as base32" >&2
		return 1
	fi

	kstore add "$key" "$secret" "$_KOTP_PROP" > /dev/null 2>&1
	if [ $? -ne 0 ]; then
		kstore update "$key" "$secret" "$_KOTP_PROP" || return 1
	fi

	echo "stored: $name"
}

function _kotp-secret {
	if [ -z "$1" ]; then
		__func_head "NAME"
		return 1
	fi

	kstore get "$(_kotp-key "$1")" "$_KOTP_PROP"
}

function _kotp-gen {
	if [ -z "$1" ]; then
		__func_head "NAME [TIME]"
		return 1
	fi

	local name when secret
	name="$1"
	when="$2"

	secret="$(_kotp-secret "$name")" || return 1

	if [ -n "$when" ]; then
		oathtool --totp --base32 --now "$when" "$secret"
	else
		oathtool --totp --base32 "$secret"
	fi
}

function _kotp-spread {
	if [ -z "$1" ]; then
		__func_head "NAME [COUNT] [TIME]"
		return 1
	fi

	local name count when secret i offset label
	name="$1"
	count="${2:-1}"
	when="$3"

	if ! [[ "$count" =~ ^[0-9]+$ ]]; then
		echo "COUNT must be a number" >&2
		return 1
	fi

	secret="$(_kotp-secret "$name")" || return 1

	for (( i=-count; i<=count; i++ )); do
		offset=$(( i * 30 ))

		if [ "$i" -eq 0 ]; then
			label="now"
		elif [ "$i" -lt 0 ]; then
			label="${offset}s"
		else
			label="+${offset}s"
		fi

		if [ -n "$when" ]; then
			printf '%8s  ' "$label"
			oathtool --totp --base32 --now "$when $offset seconds" "$secret"
		else
			printf '%8s  ' "$label"
			oathtool --totp --base32 --now "$offset seconds" "$secret"
		fi
	done
}

function _kotp-list {
	local term="$1"

	if [ -n "$term" ]; then
		kstore search "${_KOTP_PREFIX}:$term" "$_KOTP_PROP"
	else
		kstore search "${_KOTP_PREFIX}:" "$_KOTP_PROP"
	fi
}

function _kotp-del {
	if [ -z "$1" ]; then
		__func_head "NAME"
		return 1
	fi

	kstore del "$(_kotp-key "$1")" "$_KOTP_PROP"
}

function _kotp {
	local cur scope
	cur="${COMP_WORDS[COMP_CWORD]}"
	scope="${COMP_WORDS[COMP_CWORD-1]}"

	COMPREPLY=()

	case "$scope" in
		kotp)
			COMPREPLY=( $(compgen -W "set gen get spread list del" -- "$cur") )
			return
			;;
		gen|get|spread|del|delete|rm)
			COMPREPLY=( $(compgen -W "$(
				kstore search "${_KOTP_PREFIX}:" "$_KOTP_PROP" 2>/dev/null \
					| awk 'NR > 2 {print $1}' \
					| sed "s/^${_KOTP_PREFIX}://"
			)" -- "$cur") )
			return
			;;
		set)
			COMPREPLY=( $(compgen -f -- "$cur") )
			return
			;;
	esac
}

complete -F _kotp kotp