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

SELF="$(readlink -f "${BASH_SOURCE[0]}")"
export PATH="${SELF%/*}:$PATH"

cmd() {
	echo "[#] $*" >&2
	"$@"
}

auto_su() {
	[[ $UID == 0 ]] || exec sudo -p "$PROGRAM must be run as root. Please enter the password for %u to continue: " "$SELF" "${ARGS[@]}"
}

unwind() {
	set +e
	[[ -n $INTERFACE && -n $(ip link show dev "$INTERFACE" type wireguard 2>/dev/null) ]] && del_if
	exit
}

add_if() {
	ip link delete dev "$INTERFACE" 2>/dev/null || true
	cmd ip link add "$INTERFACE" type wireguard
}

del_if() {
	[[ -n $(ip link show dev "$INTERFACE" type wireguard 2>/dev/null) ]] || { echo "$PROGRAM: \`$INTERFACE' is not a WireGuard interface" >&2; exit 1; }
	if [[ $(ip route show table all) =~ .*\ dev\ $INTERFACE\ table\ ([0-9]+)\ .* ]]; then
		cmd ip rule delete table ${BASH_REMATCH[1]}
	fi
	cmd ip link delete dev "$INTERFACE"
}

up_if() {
	cmd ip link set "$INTERFACE" up
}

add_addr() {
	cmd ip address add "$1" dev "$INTERFACE"
}

add_route() {
	if [[ $1 == 0.0.0.0/0 || $1 == ::/0 ]]; then
		add_default "$1"
	else
		cmd ip route add "$1" dev "$INTERFACE"
	fi
}

add_default() {
	[[ $(join <(wg show "$INTERFACE" allowed-ips) <(wg show "$INTERFACE" endpoints)) =~ .*\ ${1//./\\.}\ ([0-9.:a-f]+):[0-9]+$ ]] && local endpoint="${BASH_REMATCH[1]}"
	[[ -n $endpoint ]] || return 0
	local table=51820
	while [[ -n $(ip route show table $table) ]]; do ((table++)); done
	cmd ip route add "$1" dev "$INTERFACE" table $table
	cmd ip rule add not to "$endpoint" table $table
}

set_config() {
	if [[ -n $CONFIG_FILE_CONTENTS ]]; then
		cmd wg setconf "$INTERFACE" <(echo "$CONFIG_FILE_CONTENTS")
	else
		cmd wg setconf "$INTERFACE" "$CONFIG_FILE"
	fi
}

save_config() {
	local old_umask="$(umask)"
	umask 077
	cmd wg showconf "$INTERFACE" > "$CONFIG_FILE.tmp" || { rm -f "$CONFIG_FILE.tmp"; exit 1; }
	mv "$CONFIG_FILE.tmp" "$CONFIG_FILE" || { rm -f "$CONFIG_FILE.tmp"; exit 1; }
	umask "$old_umask"
}

cmd_usage() {
	cat >&2 <<-_EOF
	Usage: $PROGRAM [ add | del ] INTERFACE [arguments...]

	  $PROGRAM add INTERFACE --config=CONFIG_FILE [--address=ADDRESS/CIDR...]
	               [--route=ROUTE/CIDR...] [--no-auto-route-from-allowed-ips]
	               [--env-file=ENV_FILE]

	    The add subcommand adds a new WireGuard interface, INTERFACE, replacing
	    any existing interfaces of the same name. The --config argument is
	    required, and its argument is passed to wg(8)'s setconf subcommand. The
	    --address argument(s) is recommended for this utility to be useful. The
	    --route argument is purely optional, as by default this utility will
	    automatically add routes implied by --address and as implied by the
	    allowed-ip entries inside the --config file. To disable this automatic
	    route adding, you may use the option entitled --no-auto-route-from-allowed-ips.

	  $PROGRAM del INTERFACE [--config=CONFIG_FILE_TO_SAVE] [--env-file=ENV_FILE]

	    The del subcommand removes an existing WireGuard interface. If the
	    optional --config is specified, then the existing configuration is
	    written out to the file specified, via wg(8)'s showconf subcommand.

	  $PROGRAM help

	    Show this message.

	Both \`add' and ``del' take the --env-file=ENV_FILE option. If specified,
	the contents of ENV_FILE are imported into $PROGRAM. This can be used to
	set variables in a file, instead of needing to pass them on the command
	line. The following table shows the relation between the command line
	options described above, and variables that may be declared in ENV_FILE:

	  --address=A, --address=B, --address=C       ADDRESSES=( "A" "B" "C" )
	  --route=A, --route=B, --route=C             ADDITIONAL_ROUTES=( "A" "B" "C" )
	  --config-file=F                             CONFIG_FILE="F"
	  echo C > /tmp/F, --config-file=/tmp/F       CONFIG_FILE_CONTENTS="C"
	  --no-auto-route-from-allowed-ips            AUTO_ROUTE=0

	Additionally, ENV_FILE may define the bash functions pre_add, post_add,
	pre_del, and post_del, which will be called at their respective times.
	_EOF
}

cmd_add() {
	local i
	[[ -n $CONFIG_FILE || -n $CONFIG_FILE_CONTENTS ]] || { echo "$PROGRAM: --config is required for add subcommand" >&2; exit 1; }
	auto_su
	trap unwind INT TERM EXIT
	[[ $(type -t pre_add) != function ]] || pre_add
	add_if
	set_config
	for i in "${ADDRESSES[@]}"; do
		add_addr "$i"
	done
	up_if
	if [[ $AUTO_ROUTE -eq 1 ]]; then
		for i in $(wg show "$INTERFACE" allowed-ips | grep -Po '(?<=[\t ])[0-9.:/a-f]+' | sort -nr -k 2 -t /); do
			[[ $(ip route get "$i" 2>/dev/null) == *dev\ $INTERFACE\ * ]] || add_route "$i"
		done
	fi
	for i in "${ADDITIONAL_ROUTES[@]}"; do
		add_route "$i"
	done
	[[ $(type -t post_add) != function ]] || post_add
	trap - INT TERM EXIT
}

cmd_del() {
	auto_su
	[[ $(type -t pre_del) != function ]] || pre_del
	[[ -n $CONFIG_FILE ]] && save_config
	del_if
        [[ $(type -t post_del) != function ]] || post_del
}

declare INTERFACE="$2"
declare SUBCOMMAND="$1"
declare -a ADDRESSES
declare -a ADDITIONAL_ROUTES
declare AUTO_ROUTE=1
declare CONFIG_FILE
declare CONFIG_FILE_CONTENTS
declare PROGRAM="${0##*/}"
declare -a ARGS=( "$@" )

[[ -n $INTERFACE && -n $SUBCOMMAND ]] || { cmd_usage; exit 1; }

shift 2

for arg; do
	case "$arg" in
	--env-file=*) source "${arg#*=}" ;;
	--config=*) CONFIG_FILE="${arg#*=}" ;;
	--address=*) ADDRESSES+=( ${arg#*=} ) ;;
	--route=*) ADDITIONAL_ROUTES+=( ${arg#*=} ) ;;
	--no-auto-route-from-allowed-ips) AUTO_ROUTE=0 ;;
	*) cmd_usage; exit 1 ;;
	esac
done

case "$SUBCOMMAND" in
add) cmd_add ;;
del) cmd_del ;;
*) cmd_usage; exit 1 ;;
esac

exit 0
