#!/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_taddr=
def_taddr6=
for addr in "${tap_addrs[@]}"; do
	addr="${addr%%/*}"
	if [[ "${addr}" == *:* ]]; then
		[[ -z "${def_taddr6}" ]] && def_taddr6="${addr%:1}:2" &&
			[[ -n "${def_taddr}" ]] && break
	else
		[[ -z "${def_taddr}" ]] && def_taddr="${addr%.1}.2" &&
			[[ -n "${def_taddr6}" ]] && 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
		printf -v pre '%s\n' 'chain prerouting {' \
			'type nat hook prerouting priority dstnat' \
			'policy accept'
		for daddr in "${dnat_rules[@]}"; do
			saddr= sport= dport= proto= taddr= tport=

			if [[ "${daddr}" =~ ^((.*)\ +)?to\ +(.*) ]]; then
				daddr="${BASH_REMATCH[2]}"
				taddr="${BASH_REMATCH[3]}"
			fi
			if [[ "${daddr}" =~ ^\ *from\ +(.*)\ +for\ +(.*) ||
					"${daddr}" =~ ^\ *from\ +(.*)()
					]]; then
				saddr="${BASH_REMATCH[1]}"
				daddr="${BASH_REMATCH[2]}"
			fi

			[[ "${daddr}" =~ ^\ *any\ *(/.*)?$ ]] &&
				daddr=":${BASH_REMATCH[1]}"
			ip_re='(([^:]*\[[^]]*\]([^:]*\[[^]]*\])*)?[^:]*)'
			[[ "${daddr}" =~ ^${ip_re}:?([^/]*)/?(.*) ]]
			daddr="${BASH_REMATCH[1]}"
			dport="${BASH_REMATCH[4]}"
			proto="${BASH_REMATCH[5]}"
			if [[ -n "${saddr}" ]]; then
				[[ "${saddr}" =~ ^${ip_re}:?(.*) ]]
				saddr="${BASH_REMATCH[1]}"
				sport="${BASH_REMATCH[4]}"
			fi
			if [[ -n "${taddr}" ]]; then
				[[ "${taddr}" =~ ^${ip_re}:?(.*) ]]
				taddr="${BASH_REMATCH[1]}"
				tport="${BASH_REMATCH[4]}"
			fi

			[[ "${saddr}${daddr}${taddr}" == *:* ]] &&
				ip6=1 || ip6=
			[[ -n "${sport}${dport}${tport}" ]] &&
				: "${proto:=tcp}"
			printf -v r '%s' "iifname ${wg_if}" \
				"${saddr:+ ip${ip6:+6} saddr ${saddr}}" \
				"${daddr:+ ip${ip6:+6} daddr ${daddr}}" \
				"${proto:+ meta l4proto {${proto}\}}" \
				"${sport:+ th sport ${sport}}" \
				"${dport:+ th dport ${dport}}"

			tport="${tport:+:${tport}}"
			((!ip6)) &&
				d="${taddr:-${def_taddr}}" &&
				[[ -n "${d}" ]] &&
				pre+="${r} dnat ip to ${d}${tport}"$'\n'
			((ip6)) || [[ -z "${saddr}${daddr}${taddr}" ]] &&
				d="${taddr:-[${def_taddr6}]}" &&
				[[ -n "${d}" ]] &&
				pre+="${r} dnat ip6 to ${d}${tport}"$'\n'
		done
		pre+='}'
	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}" "$@"