#!/bin/sh
# shellcheck shell=sh
#
# zt-gateway-switch — reconfigure which remote ZeroTier node acts as the
# internet exit gateway for WIBLAN clients (10.11.13.0/24).
#
# Usage:
#   zt-gateway-switch <new_ip> <mode> [drain_timeout]
#
# Modes:
#   force      - Instant cutover + conntrack flush
#   graceful   - fwmark drain; falls back to force on timeout
#
# This script is callable directly from SSH for testing; the LuCI rpcd
# backend (/usr/share/rpcd/ucode/zt-gateway.uc) wraps it for the UI.
#
# Exit codes:
#   0  success
#   1  usage error
#   2  pre-flight failed (gateway unreachable)
#   3  runtime failure (route/iptables/uci update failed)
#
# Environment overrides (used by the Docker test harness as much as by
# production):
#   ZTG_BRIDGE           bridge device (default: br-zt)
#   ZTG_WIBLAN_CIDR      WIBLAN source subnet (default: 10.11.13.0/24)
#   ZTG_TABLE_MAIN       main policy table (default: 100)
#   ZTG_TABLE_DRAIN      drain policy table (default: 101)
#   ZTG_TABLE_MWAN       mwan3 return-traffic table (default: 1)
#   ZTG_FWMARK           conntrack mark used for drain (default: 0x100)
#   ZTG_DRAIN_PRIORITY   priority of the fwmark drain rule (default: 99)
#   ZTG_HOTPLUG          hotplug script path
#   ZTG_RCLOCAL          rc.local path
#   ZTG_DRAIN_PIDFILE    pidfile for the drain monitor
#   ZTG_DRAIN_NEWFILE    file recording the new gateway IP during drain
#   ZTG_DRAIN_OLDFILE    file recording the old gateway IP during drain
#   ZTG_SKIP_PERSIST     if set to 1, skip persistence (testing)
#   ZTG_PING_IFACE       interface for the pre-flight ping (default: br-zt)
#
# This script intentionally uses only POSIX sh + busybox-compatible
# utilities so the same code path runs on the router and in the Docker
# openwrt/rootfs test image.

set -eu

# ----------------------------------------------------------------------------
# Config
# ----------------------------------------------------------------------------
BRIDGE="${ZTG_BRIDGE:-br-zt}"
WIBLAN_CIDR="${ZTG_WIBLAN_CIDR:-10.11.13.0/24}"
TABLE_MAIN="${ZTG_TABLE_MAIN:-100}"
TABLE_DRAIN="${ZTG_TABLE_DRAIN:-101}"
TABLE_MWAN="${ZTG_TABLE_MWAN:-1}"
DRAIN_PRIORITY="${ZTG_DRAIN_PRIORITY:-99}"
FWMARK="${ZTG_FWMARK:-0x100}"
HOTPLUG="${ZTG_HOTPLUG:-/etc/hotplug.d/net/99-zerotier-bridge}"
RCLOCAL="${ZTG_RCLOCAL:-/etc/rc.local}"
DRAIN_PIDFILE="${ZTG_DRAIN_PIDFILE:-/var/run/zt-gateway-drain.pid}"
DRAIN_NEWFILE="${ZTG_DRAIN_NEWFILE:-/var/run/zt-gateway-drain.new}"
DRAIN_OLDFILE="${ZTG_DRAIN_OLDFILE:-/var/run/zt-gateway-drain.old}"
PING_IFACE="${ZTG_PING_IFACE:-$BRIDGE}"
SKIP_PERSIST="${ZTG_SKIP_PERSIST:-0}"

# ----------------------------------------------------------------------------
# Logging
# ----------------------------------------------------------------------------
log() { printf '[zt-gateway-switch] %s\n' "$*" >&2; }
die() { rc=$1; shift; log "ERROR: $*"; exit "$rc"; }

# ----------------------------------------------------------------------------
# Helpers
# ----------------------------------------------------------------------------
valid_ip() {
	ip=$1
	# POSIX-only IPv4 validator (no awk regex extensions required).
	rest=$ip
	i=0
	while [ $i -lt 4 ]; do
		oct=${rest%%.*}
		if [ "$oct" = "$rest" ]; then
			# last octet
			rest=
		else
			rest=${rest#*.}
		fi
		case "$oct" in
			''|*[!0-9]*) return 1 ;;
		esac
		[ "$oct" -ge 0 ] 2>/dev/null || return 1
		[ "$oct" -le 255 ] || return 1
		i=$((i+1))
	done
	[ $i -eq 4 ] || return 1
	[ -z "$rest" ] || return 1
	[ "$ip" != "$rest" ] || return 1
	# Reject leading/trailing dots already handled above; accept.
	return 0
}

preflight_ping() {
	gw_ip=$1
	log "pre-flight: ping ${PING_IFACE}->${gw_ip}"
	if ping -c 2 -W 3 -I "$PING_IFACE" "$gw_ip" >/dev/null 2>&1; then
		return 0
	fi
	# Retry without -I in case the harness lacks the iface binding.
	if ping -c 2 -W 3 "$gw_ip" >/dev/null 2>&1; then
		return 0
	fi
	return 1
}

# Replace OLD_IP with NEW_IP in a file. IPv4 only; dots escaped for sed.
_replace_ip_in_file() {
	file=$1
	old=$2
	new=$3
	[ -f "$file" ] || { log "warning: $file not found; skipping"; return 0; }
	old_re=$(printf '%s\n' "$old" | sed 's/[.]/\\./g')
	sed -i "s/${old_re}/${new}/g" "$file"
}

set_host_route() {
	gw_ip=$1
	ip route replace "$gw_ip" dev "$BRIDGE"
}

set_table_default() {
	gw_ip=$1
	table=$2
	ip route replace default via "$gw_ip" dev "$BRIDGE" table "$table"
}

del_table_default() {
	table=$1
	ip route del default dev "$BRIDGE" table "$table" 2>/dev/null || true
}

ensure_mwan_return() {
	ip route replace "$WIBLAN_CIDR" dev "$BRIDGE" table "$TABLE_MWAN"
}

# Extract the gateway IP from the table N default route, or echo empty.
current_gateway_of() {
	table=$1
	ip route show table "$table" 2>/dev/null \
		| awk '/^[[:space:]]*default/ {
			for (i = 1; i <= NF; i++)
				if ($i == "via") { print $(i+1); exit }
		}'
}

# ----------------------------------------------------------------------------
# Persistence
# ----------------------------------------------------------------------------
persist_all() {
	new_ip=$1
	old_ip=${2:-}

	if [ "$SKIP_PERSIST" = "1" ]; then
		log "ZTG_SKIP_PERSIST=1; skipping persistence"
		return 0
	fi

	if [ -n "$old_ip" ]; then
		_replace_ip_in_file "$HOTPLUG" "$old_ip" "$new_ip"
		_replace_ip_in_file "$RCLOCAL" "$old_ip" "$new_ip"
	fi

	if command -v uci >/dev/null 2>&1; then
		if uci -q get network.zt_gateway_host >/dev/null 2>&1; then
			uci set network.zt_gateway_host.target="$new_ip"
		else
			uci -q set network.zt_gateway_host=config route
			uci -q set network.zt_gateway_host.target="$new_ip"
			uci -q set network.zt_gateway_host.interface='br-zt'
		fi

		if uci -q get network.zt_gateway_default >/dev/null 2>&1; then
			uci set network.zt_gateway_default.gateway="$new_ip"
		else
			uci -q set network.zt_gateway_default=config route
			uci -q set network.zt_gateway_default.target='0.0.0.0'
			uci -q set network.zt_gateway_default.netmask='0.0.0.0'
			uci -q set network.zt_gateway_default.gateway="$new_ip"
			uci -q set network.zt_gateway_default.table="$TABLE_MAIN"
		fi

		uci -q commit network || true
	fi
}

# ----------------------------------------------------------------------------
# Force switch
# ----------------------------------------------------------------------------
do_force() {
	new_ip=$1
	old_ip=${2:-}

	preflight_ping "$new_ip" || die 2 "gateway ${new_ip} is unreachable over ${PING_IFACE}"

	# Ensure WIBLAN traffic uses the policy table (may be missing after reboot)
	ip rule del from "$WIBLAN_CIDR" table "$TABLE_MAIN" priority 100 2>/dev/null || true
	ip rule add from "$WIBLAN_CIDR" table "$TABLE_MAIN" priority 100

	set_host_route "$new_ip"
	set_table_default "$new_ip" "$TABLE_MAIN"
	ensure_mwan_return

	if command -v conntrack >/dev/null 2>&1; then
		conntrack -D -s "$WIBLAN_CIDR" 2>/dev/null || true
	else
		log "warning: conntrack not present; skipping flush"
	fi

	persist_all "$new_ip" "$old_ip"

	if ip route show table "$TABLE_MAIN" 2>/dev/null | grep -q "default via ${new_ip}"; then
		log "force: table ${TABLE_MAIN} default via ${new_ip} confirmed"
	else
		die 3 "verification failed: default via ${new_ip} not in table ${TABLE_MAIN}"
	fi

	printf 'force switch to %s complete\n' "$new_ip"
}

# ----------------------------------------------------------------------------
# Graceful drain switch
# ----------------------------------------------------------------------------
mangle_rules_install() {
	# Source-only match: WIBLAN clients reach the router on whatever LAN
	# interface they live on; binding to -i br-zt would miss the actual
	# inbound path. We scope by source subnet so the rule fires regardless
	# of ingress interface, then CONNMARK restores/marks/saves per-flow.
	iptables -t mangle -A PREROUTING -s "$WIBLAN_CIDR" \
		-m conntrack --ctstate ESTABLISHED,RELATED -j CONNMARK --restore-mark
	iptables -t mangle -A PREROUTING -s "$WIBLAN_CIDR" \
		-m conntrack --ctstate ESTABLISHED,RELATED -m mark --mark 0 -j MARK --set-mark "$FWMARK"
	iptables -t mangle -A PREROUTING -s "$WIBLAN_CIDR" \
		-j CONNMARK --save-mark
}

mangle_rules_remove() {
	iptables -t mangle -D PREROUTING -s "$WIBLAN_CIDR" \
		-j CONNMARK --save-mark 2>/dev/null || true
	iptables -t mangle -D PREROUTING -s "$WIBLAN_CIDR" \
		-m conntrack --ctstate ESTABLISHED,RELATED -m mark --mark 0 -j MARK --set-mark "$FWMARK" 2>/dev/null || true
	iptables -t mangle -D PREROUTING -s "$WIBLAN_CIDR" \
		-m conntrack --ctstate ESTABLISHED,RELATED -j CONNMARK --restore-mark 2>/dev/null || true
}

drain_install_rules() {
	old_ip=$1
	ip route replace default via "$old_ip" dev "$BRIDGE" table "$TABLE_DRAIN"
	mangle_rules_install
	ip rule del fwmark "$FWMARK" table "$TABLE_DRAIN" priority "$DRAIN_PRIORITY" 2>/dev/null || true
	ip rule add fwmark "$FWMARK" table "$TABLE_DRAIN" priority "$DRAIN_PRIORITY"
}

drain_cleanup() {
	log "drain: cleaning up mangle rules, drain table, and pidfiles"
	mangle_rules_remove
	ip rule del fwmark "$FWMARK" table "$TABLE_DRAIN" priority "$DRAIN_PRIORITY" 2>/dev/null || true
	del_table_default "$TABLE_DRAIN"
	rm -f "$DRAIN_PIDFILE" "$DRAIN_NEWFILE" "$DRAIN_OLDFILE" 2>/dev/null || true
}

drain_count_marked() {
	# Count conntrack entries carrying the drain mark. /proc/net/nf_conntrack
	# reports marks in decimal (e.g. mark=256 for 0x100), so convert FWMARK.
	mark_dec=$(( FWMARK + 0 ))
	mark_hex="0x$(printf '%x' "$mark_dec")"

	# Preferred: conntrack CLI with -m filter (conntrack-tools >= 1.4.4).
	if command -v conntrack >/dev/null 2>&1; then
		count=$(conntrack -L -m "$mark_hex" 2>/dev/null | grep -c . || true)
		if [ -n "$count" ] && [ "$count" -gt 0 ]; then
			echo "$count"
			return
		fi
	fi

	# Fallback: parse /proc/net/nf_conntrack. Works without the CLI and
	# also catches entries the CLI filter might miss on older builds.
	awk -v mdec="$mark_dec" -v mhex="$mark_hex" '
		BEGIN { n = 0 }
		{
			if (index($0, "mark=" mdec) || index($0, "mark=" mhex)) n++
		}
		END { print n }
	' /proc/net/nf_conntrack 2>/dev/null || echo 0
}

drain_monitor() {
	new_ip=$1
	old_ip=$2
	timeout=$3

	# Grace period before the first count: packets already in flight need
	# at least one PREROUTING pass through the mangle rules before their
	# conntrack entry picks up mark=0x100. Without this cushion the very
	# first count races the mangle hook and reports zero prematurely.
	DRAIN_INITIAL_GRACE="${ZTG_DRAIN_INITIAL_GRACE:-3}"
	# Number of consecutive zero-count polls required before we declare
	# the drain complete. Guards against transient empty reads that race
	# the conntrack update.
	DRAIN_MIN_ZERO_POLLS="${ZTG_DRAIN_MIN_ZERO_POLLS:-2}"
	DRAIN_POLL_INTERVAL="${ZTG_DRAIN_POLL_INTERVAL:-5}"

	sleep "$DRAIN_INITIAL_GRACE"

	end=$(( $(date +%s) + timeout ))
	consecutive_zeros=0
	while :; do
		remaining=$(drain_count_marked)
		log "drain: ${remaining} marked connection(s) remaining (consecutive_zeros=${consecutive_zeros})"
		if [ "$remaining" -eq 0 ]; then
			consecutive_zeros=$(( consecutive_zeros + 1 ))
			if [ "$consecutive_zeros" -ge "$DRAIN_MIN_ZERO_POLLS" ]; then
				log "drain: complete (${DRAIN_MIN_ZERO_POLLS} consecutive zero polls)"
				drain_cleanup
				return 0
			fi
		else
			consecutive_zeros=0
		fi
		if [ "$(date +%s)" -ge "$end" ]; then
			log "drain: timeout (${timeout}s) reached; falling back to force"
			set_host_route "$new_ip"
			set_table_default "$new_ip" "$TABLE_MAIN"
			ensure_mwan_return
			if command -v conntrack >/dev/null 2>&1; then
				conntrack -D -s "$WIBLAN_CIDR" 2>/dev/null || true
			fi
			drain_cleanup
			persist_all "$new_ip" "$old_ip"
			return 0
		fi
		sleep "$DRAIN_POLL_INTERVAL"
	done
}

do_graceful() {
	new_ip=$1
	timeout=${2:-600}

	old_ip=$(current_gateway_of "$TABLE_MAIN")
	if [ -z "$old_ip" ]; then
		log "graceful: no current gateway in table ${TABLE_MAIN}; falling back to force"
		do_force "$new_ip" ""
		return $?
	fi

	if [ "$old_ip" = "$new_ip" ]; then
		log "graceful: new gateway equals current gateway (${new_ip}); nothing to do"
		return 0
	fi

	preflight_ping "$new_ip" || die 2 "gateway ${new_ip} is unreachable over ${PING_IFACE}"
	# Ensure WIBLAN traffic uses the policy table (may be missing after reboot)
	ip rule del from "$WIBLAN_CIDR" table "$TABLE_MAIN" priority 100 2>/dev/null || true
	ip rule add from "$WIBLAN_CIDR" table "$TABLE_MAIN" priority 100


	printf '%s\n' "$new_ip" >"$DRAIN_NEWFILE"
	printf '%s\n' "$old_ip" >"$DRAIN_OLDFILE"

	drain_install_rules "$old_ip"

	set_host_route "$new_ip"
	set_table_default "$new_ip" "$TABLE_MAIN"
	ensure_mwan_return

	persist_all "$new_ip" "$old_ip"

	: >"$DRAIN_PIDFILE"
	(
		trap 'drain_cleanup; exit 0' TERM INT
		echo $$ >"$DRAIN_PIDFILE"
		drain_monitor "$new_ip" "$old_ip" "$timeout"
		rm -f "$DRAIN_PIDFILE" "$DRAIN_NEWFILE" "$DRAIN_OLDFILE" 2>/dev/null || true
	) >/tmp/zt-gateway-drain.log 2>&1 &

	printf 'graceful switch to %s started (draining from %s, timeout %ss)\n' \
		"$new_ip" "$old_ip" "$timeout"
}

# ----------------------------------------------------------------------------
# Entry point
# ----------------------------------------------------------------------------
main() {
	if [ $# -lt 2 ]; then
		cat >&2 <<'USAGE'
Usage: zt-gateway-switch <new_ip> <mode> [drain_timeout]

Modes:
  force      Instant cutover + conntrack flush
  graceful   fwmark drain; falls back to force on timeout

Options:
  drain_timeout   graceful-mode drain timeout in seconds (default: 600)
USAGE
		die 1 "missing arguments"
	fi

	new_ip=$1
	mode=$2
	timeout=${3:-600}

	if ! valid_ip "$new_ip"; then
		die 1 "invalid IPv4 address: ${new_ip}"
	fi

	case "$mode" in
		force)
			old_ip=$(current_gateway_of "$TABLE_MAIN")
			do_force "$new_ip" "${old_ip:-}"
			;;
		graceful)
			do_graceful "$new_ip" "$timeout"
			;;
		*)
			die 1 "unknown mode: ${mode} (use force|graceful)"
			;;
	esac
}

main "$@"