#!/usr/bin/env bash
#
# Copyright 2026 David Vazgenovich Shakaryan

die() {
	printf '%s\n' "${0##*/}: $*" >&2
	exit 1
}

[[ -n "${SUDO_USER}" ]] || die 'SUDO_USER is not set'

run() {
	if [[ -t 1 ]]; then
		printf '\e[90m[*] %s\e[0m\n' "$*"
	else
		printf '[*] %s\n' "$*"
	fi

	"$@"
}

wg_conf=
netns=
enable_tap=
tap_addrs=()
dnat_rules=()
while (($#)); do
	case "$1" in
		-c|--config)
			wg_conf="${2}"
			shift 2 ;;
		-n|--netns)
			netns="${2}"
			shift 2 ;;
		--tap)
			enable_tap=1
			shift ;;
		--tap-addr)
			tap_addrs+=("${2}")
			shift 2 ;;
		--dnat)
			dnat_rules+=("${2}")
			shift 2 ;;
		--)
			shift
			break ;;
		-*)
			die "invalid option: ${1}" ;;
		*)
			break ;;
	esac
done

: ${wg_conf:=default}
[[ "${wg_conf}" =~ ^[A-Za-z0-9+.=_-]{1,15}$ ]] &&
	wg_conf="/etc/wg-netns-exec/${wg_conf}.conf"
[[ -e "${wg_conf}" ]] || die "config file does not exist: ${wg_conf}"
echo "config: $wg_conf"
: ${netns:=$(LC_ALL=C tr -dc 'a-z0-9' < /dev/urandom | head -c 8)}
wg_if="wg-${netns}"

((!${#tap_addrs[@]})) && tap_addrs=(10.0.0.1/32 fd00::1/64)
def_daddr=
def_daddr6=
for addr in "${tap_addrs[@]}"; do
	addr="${addr%%/*}"
	if [[ "${addr}" == *:* ]]; then
		[[ -z "${def_daddr6}" ]] && def_daddr6="${addr%:1}:2" &&
			[[ -n "${def_daddr}" ]] && break
	else
		[[ -z "${def_daddr}" ]] && def_daddr="${addr%.1}.2" &&
			[[ -n "${def_daddr6}" ]] && break
	fi
done

run ip netns del "${netns}" 2>/dev/null
run ip netns add "${netns}"
trap 'run ip netns del "${netns}"' EXIT

run ip link add "${wg_if}" type wireguard
run ip link set "${wg_if}" netns "${netns}"
for i in $(awk '
			/^\[Interface\]$/ { m = 1; next } /^\[/ { m = 0 }
			m && sub(/^Address\s*=\s*/, "") {
				gsub(/,\s*/, "\n"); print }
		' "${wg_conf}"); do
	run ip -n "${netns}" addr add "${i}" dev "${wg_if}"
done
run ip netns exec "${netns}" wg setconf "${wg_if}" <(
	wg-quick strip "${wg_conf}")
run ip -n "${netns}" link set "${wg_if}" up
run ip -n "${netns}" -4 route add default dev "${wg_if}"
run ip -n "${netns}" -6 route add default dev "${wg_if}"
run ip netns exec "${netns}" curl -s4 -m 10 ifconfig.co
run ip netns exec "${netns}" curl -s6 -m 10 ifconfig.co

if ((enable_tap)); then
	run ip netns exec "${netns}" ip tuntap add dev tap0 mode tap user \
		"${SUDO_USER}"
	run ip netns exec "${netns}" ip link set tap0 up
	for addr in "${tap_addrs[@]}"; do
		run ip netns exec "${netns}" ip addr add "${addr}" dev tap0
	done

	run ip netns exec "${netns}" sysctl -wq net.ipv4.conf.all.forwarding=1
	run ip netns exec "${netns}" sysctl -wq net.ipv6.conf.all.forwarding=1

	pre=
	if ((${#dnat_rules[@]})); then
		pre='chain prerouting {
			type nat hook prerouting priority dstnat
			policy accept'
		for saddr in "${dnat_rules[@]}"; do
			sport= proto= daddr= dport=
			if [[ "${saddr}" =~ ^((.*)\ +)?to\ +(.*) ]]; then
				saddr="${BASH_REMATCH[2]}"
				daddr="${BASH_REMATCH[3]}"
			fi

			ip_re='(([^:]*\[[^]]*\]([^:]*\[[^]]*\])*)?[^:]*)'
			[[ "${saddr}" =~ ^\ *any\ *(/.*)?$ ]] &&
				saddr=":${BASH_REMATCH[1]}"
			[[ "${saddr}" =~ ^${ip_re}:?([^/]*)/?(.*) ]]
			saddr="${BASH_REMATCH[1]}"
			sport="${BASH_REMATCH[4]}"
			proto="${BASH_REMATCH[5]}"

			if [[ -n "${daddr}" ]]; then
				[[ "${daddr}" =~ ^${ip_re}:?(.*) ]]
				daddr="${BASH_REMATCH[1]}"
				dport="${BASH_REMATCH[4]}"
			fi

			[[ "${saddr}" == *:* || "${daddr}" == *:* ]] &&
				ip6=1 || ip6=

			r="iifname ${wg_if}"
			[[ -n "${saddr}" ]] && r+=" ip${ip6:+6} saddr ${saddr}"
			[[ -n "${sport}" || -n "${dport}" ||
					-n "${proto}" ]] &&
				r+=" meta l4proto {${proto:-tcp}}"
			[[ -n "${sport}" ]] && r+=" th dport ${sport}"
			dport="${dport:+:${dport}}"

			((!ip6)) &&
				d="${daddr:-${def_daddr}}" &&
				[[ -n "${d}" ]] &&
				pre+=$'\n'"${r} dnat ip to ${d}${dport}"

			((ip6)) || [[ -z "${saddr}" && -z "${daddr}" ]] &&
				d="${daddr:-${def_daddr6}}" &&
				[[ -n "${d}" ]] &&
				pre+=$'\n'"${r} dnat ip6 to ${d}${dport}"
		done
		pre+=$'\n}'
	fi
	run ip netns exec "${netns}" nft -f - <<-EOF
		table inet nat {
			${pre}
			chain postrouting {
				type nat hook postrouting priority srcnat
				policy accept
				oifname ${wg_if} masquerade
			}
		}
	EOF
	run ip netns exec "${netns}" nft list ruleset
fi

(($#)) || set -- sleep infinity
run ip netns exec "${netns}" sudo -u "${SUDO_USER}" "$@"