#!/usr/bin/env mksh
# Script by Ypnose - http://ywstd.fr
set -e

WG_TOOL="${0##*/}"
WG_CONF="$1"
WG_MARK="9999"
WG_FMON="/tmp/vpn_${USER_ID}"

########################################################################

function p_str {
	builtin print -r -- "$1"
}

function p_mes {
	p_str "[${WG_TOOL}] $1"
}

function p_err {
	builtin print -ru2 -- "$1"
	exit 1
}

function split_wg_value {
	typeset IFS v f x
	# Use nameref so we can use "shortcut" variable to array
	typeset -n t="$1"

	IFS=','
	for v in $2; do
		f=0
		for x in "${t[@]}"; do
			if [[ $v = "$x" ]]; then f=1; fi
		done
		if (( f == 0 )); then t+=( "$v" ); fi
	done
}

function read_wg_conf {
	typeset line WG_ALLOW_NET_PARAM WG_STRIP_OPT WG_FIELD WG_VALUE \
		OLD_WG_STRIP_OPT
	set -A WG_IPADDR
	set -A WG_DNS

	WG_INF="${WG_CONF##*/}"
	WG_INF="${WG_INF%.conf}"

	while IFS= read -r line; do
		# If line is empty or a comment, ignore it
		[[ -z $line || $line = *#* ]] && continue
		# When Interface section starts, allow local network parameters
		[[ $line = "[Interface]" ]] && WG_ALLOW_NET_PARAM=1
		# When Peer section starts, forbid local network parameters
		[[ $line = "[Peer]" ]] && WG_ALLOW_NET_PARAM=0
		WG_STRIP_OPT="${line// /}"
		# Avoid duplicate lines
		[[ $WG_STRIP_OPT = "$OLD_WG_STRIP_OPT" ]] && continue
		WG_FIELD="${WG_STRIP_OPT%%=*}"
		WG_VALUE="${WG_STRIP_OPT#*=}"
		OLD_WG_STRIP_OPT="$WG_STRIP_OPT"
		# Record local port to open. Only first value is matched
		if [[ $WG_FIELD = "Endpoint" && -z $WG_PORT ]]; then
			if [[ ${WG_VALUE##*:} != "$WG_VALUE" ]]; then
				WG_PORT="${WG_VALUE##*:}"
			fi
		fi
		# Create stripped config for wg
		if [[ $WG_FIELD = @(\[Interface\]|\[Peer\]|AllowedIPs|Endpoint|ListenPort|PersistentKeepalive|@(Preshared|Private|Public)Key) ]]; then
			WG_CONF_STDIN+=( "$line" )
			continue
		fi
		[[ $WG_FIELD = "$WG_VALUE" || -z $WG_VALUE ]] && continue
		# Catch and "record" network parameters
		if (( WG_ALLOW_NET_PARAM == 1 )); then
			case $WG_FIELD in
				Address)
					split_wg_value WG_IPADDR "$WG_VALUE"
					continue
					;;
				DNS)
					split_wg_value WG_DNS "$WG_VALUE"
					continue
					;;
				MTU)
					# Only first MTU value is recorded
					[[ -z $WG_MTU ]] && WG_MTU="$WG_VALUE"
					continue
					;;
			esac
		fi
	done <"$WG_CONF"
}

function print_wg_conf {
	typeset a d

	read_wg_conf
	for o in "${WG_CONF_STDIN[@]}"; do p_str "$o"; done
	p_str "[Network]"
	for a in "${WG_IPADDR[@]}"; do p_str "IP = $a"; done
	for d in "${WG_DNS[@]}"; do p_str "DNS = $d"; done
	p_str "PORT = $WG_PORT"
	if [[ -n $WG_MTU ]]; then p_str "MTU = $WG_MTU"; fi
}

function check_wg_inf {
	ip -br link show "$WG_INF" >/dev/null 2>&1
}

function verify_wg_conf {
	typeset a

	# Apply read-write only for the current user
	chmod 600 "$WG_CONF"
	# Make sure previous interface was cleaned up
	if check_wg_inf; then
		p_err "WireGuard $WG_INF interface is already created"
	fi
	# Make sure netmask is defined
	for a in "${WG_IPADDR[@]}"; do
		if [[ $a = "${a#*/}" ]]; then
			p_err "$a address in $WG_CONF is invalid (undefined netmask)"
		fi
	done
	# If remote port was not defined in the config file, that's an issue
	if [[ -z $WG_PORT ]]; then
		p_err "Endpoint port is not defined in $WG_CONF"
	fi
	# If MTU is > 1500, that's an issue
	if (( WG_MTU > 1500 )); then
		p_err "Invalid MTU value in $WG_CONF (too high)"
	fi
}

function move_resolv {
	if [[ -e $1 ]]; then mv "$1" "$2"; fi && chmod 444 "$2"
}

function manage_fwport {
	iptables "$1" OUTPUT -p udp -m udp --dport "$WG_PORT" \
		-m conntrack --ctstate NEW -j ACCEPT
}

function get_remote_ip {
	curl -Ss -L https://ipinfo.io/ip
}

function set_wg_network {
	typeset o a d

	# Make sure it is launched as root
	if (( USER_ID != 0 )); then
		p_err "$WG_TOOL needs to be launched as root"
	fi

	trap '' ABRT
	trap 'stop_wg_con' HUP EXIT


	# Create WireGuard interface
	ip link add dev "$WG_INF" type wireguard
	p_mes "Created $WG_INF interface"
	# Pass stripped config to wg
	for o in "${WG_CONF_STDIN[@]}"; do p_str "$o"; done \
		| wg setconf "$WG_INF" /dev/stdin
	# Add adresses on the WireGuard interface
	for a in "${WG_IPADDR[@]}"; do
		ip address add "$a" dev "$WG_INF"
		p_mes "Added $a address to $WG_INF"
	done
	# Set DNS
	if [[ -n ${WG_DNS[0]} ]]; then
		move_resolv /etc/resolv.conf /etc/resolv.conf.before_wg
		for d in "${WG_DNS[@]}"; do
			p_str "nameserver $d" >>/etc/resolv.conf
			p_mes "Added $d nameserver to /etc/resolv.conf"
		done
	fi
	# Set MTU
	if [[ -z $WG_MTU ]]; then
		# With IPv6, 80 bytes are used, with IPv4 only, 60 bytes are used
		# https://lists.zx2c4.com/pipermail/WireGuard/2017-December/002201.html
		IFS= read -r WG_MTU <"/sys/class/net/${WG_INF}/mtu"
	fi
	ip link set mtu "$WG_MTU" dev "$WG_INF"
	ip link set dev "$WG_INF" up
	# Set default route
	# Mark packets coming out of WireGuard interface
	wg set "$WG_INF" fwmark "$WG_MARK"
	# Define table XXXX as default route
	ip route add default dev "$WG_INF" table "$WG_MARK"
	# If packet doesn't come from WireGuard interface, force it to use
	# this interface
	ip rule add not fwmark "$WG_MARK" table "$WG_MARK"
	# Prevent packets for LAN to go through the VPN
	ip rule add table main suppress_prefixlength 0
	p_mes "Default route updated"
	# Set firewall rules
	manage_fwport "-A"
	p_mes "Firewall rule added for $WG_PORT (OUTPUT)"
	# Print current IP
	#get_remote_ip
	p_mes "Current IP: $(get_remote_ip)"
	p_mes "$WG_TOOL PID: $$"
	>"${WG_FMON}"
	# Wait for the user to end the tunnel with SIGINT, SIGQUIT
	# or SIGTERM
	while : ; do sleep 86400; done
}

function start_wg_con {
	read_wg_conf
	verify_wg_conf
	set_wg_network
}

function stop_wg_con {
	set +e
	p_str
	# Delete firewall rules
	while [[ $(iptables -L OUTPUT -n) = *ACCEPT+( )udp*udp*dpt:"${WG_PORT:-0}"+( )ctstate+( )NEW ]]; do
		manage_fwport -D
	done
	p_mes "Firewall rules deleted"
	# Delete route rules
	while [[ $(ip rule) = *lookup+( )"$WG_MARK"* ]]; do
		ip rule delete table "$WG_MARK"
	done
	while [[ $(ip rule) = *lookup+( )main+( )suppress_prefixlength+( )0* ]]; do
		ip rule delete table main suppress_prefixlength 0
	done
	p_mes "Previous default route restored"
	# Restore old DNS
	if [[ -n ${WG_DNS[0]} ]]; then
		move_resolv /etc/resolv.conf.before_wg /etc/resolv.conf
		p_mes "Restored /etc/resolv.conf"
	fi
	# Flush address and WireGuard interface
	if check_wg_inf; then
		ip address flush dev "$WG_INF"
		ip link delete "$WG_INF" type wireguard
		p_mes "Deleted $WG_INF interface"
	fi
	#get_remote_ip
	[[ -e $WG_FMON ]] && rm "$WG_FMON"
	p_str "DONE"
	exit
}

########################################################################

if [[ -z $1 ]]; then
	p_err "usage: $WG_TOOL [WG_CONFIG] [-c]"
fi

if (( $# == 2 )); then
	# Display config network parameters to stdout
	if [[ $2 = "-c" ]]; then
		print_wg_conf "$1"
		exit
	fi
fi

start_wg_con "$1"