#!/usr/bin/env bash
# Unified Linux VPN killswitch for OpenVPN, WireGuard/wg-quick, and
# NetworkManager dispatcher.d. It can enforce a system-wide killswitch or a
# single-user-only one, optionally allow or block directly-connected LAN
# subnets, force plain DNS on port 53 to the VPN DNS servers, and allow only
# the VPN tunnel, loopback, LAN (if enabled), VPN endpoint, and VPN DNS while
# active. This helps prevent normal IP leaks and plain DNS leaks if the VPN
# drops.
#
# If ALLOW_LAN=1, local network access (e.g. printers, NAS) is allowed
# outside the tunnel. If ALLOW_LAN=0, all non-VPN traffic—including LAN—is
# blocked, which provides stronger isolation and helps mitigate attacks like
# TunnelCrack that rely on tricking the system into sending traffic outside
# the VPN via local or non-tunneled routes.
#
# Note: This script only enforces DNS usage at the network level (port 53).
# It does NOT block DNS-over-HTTPS (DoH) or DNS-over-TLS (DoT). Applications
# using DoH/DoT may still send DNS queries to third parties. Those queries
# should still be routed through the VPN (so your real IP is not exposed),
# but they can still be considered DNS leaks from a privacy perspective.
#
# Supported ways to call it:
#   - OpenVPN up/down script:
#       script-security 2
#       up /path/to/this.sh
#       down /path/to/this.sh
#
#   - WireGuard via wg-quick:
#       PostUp = /path/to/this.sh on %i
#       PostDown = /path/to/this.sh off %i
#     (%i expands to the interface name, e.g. wg0)
#
#   - NetworkManager dispatcher.d:
#       place or symlink it in /etc/NetworkManager/dispatcher.d/
#       It understands dispatcher events such as vpn-up/vpn-down and
#       WireGuard-style up/down events.
#
#   - Standalone usage after bringing a tunnel up yourself:
#       /path/to/this.sh on
#       /path/to/this.sh on cs-poland
#       /path/to/this.sh off
#       /path/to/this.sh status
#
# Standalone mode is mainly intended for WireGuard or for OpenVPN only when
# VPN_HOST or VPN_ENDPOINT_V4/V6 is configured. OpenVPN users should prefer
# --up/--down hooks or NetworkManager dispatcher integration.
set -Eeuo pipefail
shopt -s inherit_errexit 2>/dev/null || true

###############################################################################
# Configuration
###############################################################################

# auto | nft | iptables
BACKEND="auto"

# system | user
MODE="system"

# Used only when MODE="user"
KS_USER="justme"

# 1 = allow access to directly-connected LAN prefixes while VPN is up
# 0 = block LAN too (Tunnelcrack protection)
ALLOW_LAN=1

# VPN DNS servers pushed by the server
DNS_V4="10.31.33.7"  # change to 10.31.33.8 if you don't want ad/tracker blocking
DNS_V6="2001:db8::7" # change to 2001:db8::8 if you don't want ad/tracker blocking

# Optional fallback VPN hostname to resolve when OpenVPN didn't tell us the
# remote endpoint IP(s). Example:
# VPN_HOST="paris.cstorm.is"
VPN_HOST=""

# Optional hard overrides
VPN_ENDPOINT_V4=""
VPN_ENDPOINT_V6=""

# Optional tunnel interface override; otherwise derived from OpenVPN/NM/standalone
VPN_IF_OVERRIDE=""

# State
STATE_DIR="/run/vpn-killswitch"
STATE_FILE="$STATE_DIR/state.env"
STATE_FILE_TMP="$STATE_DIR/state.env.tmp"
RESOLV_BACKUP="$STATE_DIR/resolv.conf.backup"

# Internal safety state
ENABLE_IN_PROGRESS=0
RESOLV_CHANGED=0
ROLLBACK_ACTIVE=0
NFT_TMPFILE=""

# Policy routing for MODE=user
ROUTE_TABLE_ID=51820
RULE_PREF=51820

# nftables object names
NFT_TABLE_INET="vpnks"
NFT_TABLE_IP_NAT="vpnks_nat4"
NFT_TABLE_IP6_NAT="vpnks_nat6"

# iptables chain names
IPT_CHAIN_IN="VPNKS_IN"
IPT_CHAIN_OUT="VPNKS_OUT"
IPT_CHAIN_OUT_USER="VPNKS_OUT_USER"
IPT_CHAIN_NAT="VPNKS_DNS"

###############################################################################
# Helpers
###############################################################################

log() { printf '%s\n' "$*"; }
err() { printf 'Error: %s\n' "$*" >&2; exit 1; }
warn() { printf 'Warning: %s\n' "$*" >&2; }

need_root() {
    [[ "$(id -u)" -eq 0 ]] || err "run this script as root"
}

have() {
    command -v "$1" >/dev/null 2>&1
}

is_ipv4() {
    [[ "${1:-}" =~ ^([0-9]{1,3}\.){3}[0-9]{1,3}$ ]]
}

is_ipv6() {
    [[ "${1:-}" == *:* ]]
}

cleanup_temp_files() {
    [[ -n "${NFT_TMPFILE:-}" ]] && rm -f -- "$NFT_TMPFILE" 2>/dev/null || true
    NFT_TMPFILE=""
}

try_modprobe() {
    local mod="$1"
    have modprobe || return 0
    modprobe "$mod" >/dev/null 2>&1 || true
}

backend_usable_nft() {
    have nft || return 1
    try_modprobe nft_chain_nat
    try_modprobe nft_nat
    nft list tables >/dev/null 2>&1 || return 1
    return 0
}

backend_usable_iptables() {
    have iptables && have ip6tables || return 1
    try_modprobe iptable_nat
    try_modprobe ip6table_nat
    iptables -t nat -L -n >/dev/null 2>&1 || return 1
    ip6tables -t nat -L -n >/dev/null 2>&1 || return 1
    return 0
}

check_backend_capabilities() {
    case "${FW_BACKEND:-}" in
        nft)
            have nft || err "nft backend selected but nft command not found"

            # Best-effort module nudges for older systems
            try_modprobe nft_chain_nat
            try_modprobe nft_nat
            try_modprobe nf_tables
            ;;

        iptables)
            have iptables || err "iptables backend selected but iptables command not found"
            have ip6tables || err "iptables backend selected but ip6tables command not found"

            # Best-effort module nudges for older systems
            try_modprobe ip_tables
            try_modprobe iptable_filter
            try_modprobe iptable_nat
            try_modprobe ip6_tables
            try_modprobe ip6table_filter
            try_modprobe ip6table_nat
            try_modprobe nf_nat
            try_modprobe nf_conntrack

            iptables -t nat -L -n >/dev/null 2>&1 \
                || err "iptables IPv4 nat table is unavailable"

            ip6tables -t nat -L -n >/dev/null 2>&1 \
                || err "ip6tables IPv6 nat table is unavailable"

            iptables -L -n >/dev/null 2>&1 \
                || err "iptables filter table is unavailable"

            ip6tables -L -n >/dev/null 2>&1 \
                || err "ip6tables filter table is unavailable"
            ;;
    esac
}

wait_for_wireguard_ready() {
    local ifname="$1"
    local tries="${2:-20}"
    local endpoint=""
    local hs=""

    iface_is_wireguard "$ifname" || return 0

    while (( tries > 0 )); do
        iface_exists "$ifname" || { sleep 0.5; ((tries--)); continue; }

        endpoint="$(wg show "$ifname" endpoints 2>/dev/null | awk 'NR==1 {print $2}' || true)"
        hs="$(wg show "$ifname" latest-handshakes 2>/dev/null | awk 'NR==1 {print $2}' || true)"

        if [[ -n "$endpoint" && -n "$hs" && "$hs" != "0" ]]; then
            return 0
        fi

        sleep 0.5
        ((tries--))
    done

    warn "WireGuard interface '$ifname' did not look fully ready before timeout; continuing anyway"
    return 0
}

iface_is_openvpn_tun() {
    local ifname="${1:-}"
    [[ -n "$ifname" ]] || return 1
    [[ "$ifname" =~ ^tun[0-9A-Za-z_.-]*$ ]] || return 1
    iface_exists "$ifname"
}

rollback_enable() {
    [[ "${ROLLBACK_ACTIVE:-0}" -eq 1 ]] && return 0
    ROLLBACK_ACTIVE=1

    warn "enable failed; attempting rollback"

    if [[ "${FW_BACKEND:-}" == "nft" ]]; then
        nft_remove_all 2>/dev/null || true
    elif [[ "${FW_BACKEND:-}" == "iptables" ]]; then
        iptables_remove_all 2>/dev/null || true
    fi

    [[ "${RESOLV_CHANGED:-0}" -eq 1 ]] && restore_resolv_conf 2>/dev/null || true
    rm -f -- "$STATE_FILE_TMP" "$STATE_FILE" 2>/dev/null || true
    cleanup_temp_files
}

mkdir_state() {
    [[ -d /run ]] || err "/run does not exist"
    mkdir -p -- "$STATE_DIR" || err "failed to create state dir: $STATE_DIR"
    [[ -d "$STATE_DIR" ]] || err "state dir is not a directory: $STATE_DIR"
    chmod 700 -- "$STATE_DIR" || err "failed to chmod state dir: $STATE_DIR"
    [[ -w "$STATE_DIR" ]] || err "state dir is not writable: $STATE_DIR"
}

choose_backend() {
    case "$BACKEND" in
        nft)
            have nft || err "BACKEND=nft but nft not found"
            FW_BACKEND="nft"
            ;;
        iptables)
            have iptables || err "BACKEND=iptables but iptables not found"
            have ip6tables || err "BACKEND=iptables but ip6tables not found"
            FW_BACKEND="iptables"
            ;;
        auto)
            if backend_usable_nft; then
                FW_BACKEND="nft"
            elif backend_usable_iptables; then
                FW_BACKEND="iptables"
            else
                err "no usable firewall backend found (need working nft or iptables/ip6tables nat support)"
            fi
            ;;
        *)
            err "invalid BACKEND=$BACKEND"
            ;;
    esac
}

detect_action_and_iface() {
    ACTION=""
    VPN_IF="${VPN_IF_OVERRIDE}"

    # NetworkManager dispatcher: $1=<iface>, $2=<action>
    case "${2:-}" in
        vpn-up)
            ACTION="enable"
            [[ -n "$VPN_IF" ]] || VPN_IF="${1:-}"
            return
            ;;
        vpn-down)
            ACTION="disable"
            [[ -n "$VPN_IF" ]] || VPN_IF="${1:-}"
            return
            ;;
        up)
            if [[ -n "${1:-}" ]] && iface_is_wireguard "${1:-}"; then
                ACTION="enable"
                [[ -n "$VPN_IF" ]] || VPN_IF="${1:-}"
                return
            fi
            ;;
        down)
            if [[ -n "${1:-}" ]] && [[ -r "$STATE_FILE" ]]; then
                load_state || true
                if [[ "${VPN_IF:-}" == "${1:-}" ]]; then
                    ACTION="disable"
                    return
                fi
            fi
            ;;
    esac

    # OpenVPN script hooks
    if [[ "${script_type:-}" == "up" ]]; then
        ACTION="enable"
        [[ -n "$VPN_IF" ]] || VPN_IF="${dev:-${1:-}}"
        return
    elif [[ "${script_type:-}" == "down" ]]; then
        ACTION="disable"
        [[ -n "$VPN_IF" ]] || VPN_IF="${dev:-${1:-}}"
        return
    fi

    # Standalone
    case "${1:-}" in
        on|enable)
            ACTION="enable"
            [[ -n "$VPN_IF" ]] || VPN_IF="${2:-}"
            [[ -n "$VPN_IF" ]] || VPN_IF="$(detect_single_vpn_iface || true)"
            ;;
        off|disable)
            ACTION="disable"
            [[ -n "$VPN_IF" ]] || VPN_IF="${2:-}"
            ;;
        status)
            ACTION="status"
            ;;
        *)
            cat >&2 <<'EOF'
Usage:
  vpn-killswitch on [wg0|tun0]
  vpn-killswitch off [wg0|tun0]
  vpn-killswitch status

Notes:
  - Standalone mode works as-is for WireGuard, but for OpenVPN, setting VPN_HOST or VPN_ENDPOINT_V4/VPN_ENDPOINT_V6 is required.
  - OpenVPN users should prefer --up/--down hooks or NetworkManager dispatcher.

OpenVPN:
  up /usr/local/sbin/vpn-killswitch
  down /usr/local/sbin/vpn-killswitch
  script-security 2

NetworkManager dispatcher:
  symlink or call this script from /etc/NetworkManager/dispatcher.d/
  It understands vpn-up and vpn-down.
EOF
            exit 1
            ;;
    esac
}

resolve_vpn_endpoints() {
    ENDPOINT_V4="${VPN_ENDPOINT_V4}"
    ENDPOINT_V6="${VPN_ENDPOINT_V6}"

    # OpenVPN usually provides trusted_ip. On v6-capable setups it may contain
    # either family depending on what the control channel used.
    if [[ -n "${trusted_ip:-}" ]]; then
        if is_ipv6 "$trusted_ip"; then
            ENDPOINT_V6="$trusted_ip"
        elif is_ipv4 "$trusted_ip"; then
            ENDPOINT_V4="$trusted_ip"
        fi
    fi

    # Some setups expose trusted_ip6 separately
    if [[ -n "${trusted_ip6:-}" ]] && is_ipv6 "${trusted_ip6:-}"; then
        ENDPOINT_V6="$trusted_ip6"
    fi

    # Fallback: resolve configured VPN_HOST
    if [[ -z "$ENDPOINT_V4" && -n "$VPN_HOST" ]] && have getent; then
        ENDPOINT_V4="$(getent ahostsv4 "$VPN_HOST" | awk 'NR==1{print $1; exit}')"
    fi
    if [[ -z "$ENDPOINT_V6" && -n "$VPN_HOST" ]] && have getent; then
        ENDPOINT_V6="$(getent ahostsv6 "$VPN_HOST" | awk 'NR==1{print $1; exit}')"
    fi
    
    # WireGuard fallback: derive endpoint from the live interface if available.
    # Non-fatal on non-WireGuard interfaces.
    if [[ -n "${VPN_IF:-}" ]] && have wg; then
        local wg_ep=""
        if wg show interfaces 2>/dev/null | tr ' ' '\n' | grep -Fxq "$VPN_IF"; then
            wg_ep="$(wg show "$VPN_IF" endpoints 2>/dev/null | awk 'NR==1 {print $2}' || true)"
        fi

        if [[ -n "$wg_ep" ]]; then
            if [[ "$wg_ep" =~ ^\[([0-9a-fA-F:]+)\]:[0-9]+$ ]]; then
                [[ -z "$ENDPOINT_V6" ]] && ENDPOINT_V6="${BASH_REMATCH[1]}"
            elif [[ "$wg_ep" =~ ^([0-9]+\.[0-9]+\.[0-9]+\.[0-9]+):[0-9]+$ ]]; then
                [[ -z "$ENDPOINT_V4" ]] && ENDPOINT_V4="${BASH_REMATCH[1]}"
            fi
        fi
    fi
}

save_state() {
    mkdir_state
    cat > "$STATE_FILE_TMP" <<EOF
FW_BACKEND=$FW_BACKEND
MODE=$MODE
KS_USER=$KS_USER
ALLOW_LAN=$ALLOW_LAN
DNS_V4=$DNS_V4
DNS_V6=$DNS_V6
VPN_IF=$VPN_IF
ENDPOINT_V4=${ENDPOINT_V4:-}
ENDPOINT_V6=${ENDPOINT_V6:-}
ROUTE_TABLE_ID=$ROUTE_TABLE_ID
RULE_PREF=$RULE_PREF
EOF
    chmod 600 "$STATE_FILE_TMP" || err "failed to chmod temp state file"
    mv -f -- "$STATE_FILE_TMP" "$STATE_FILE" || err "failed to install state file"
}

load_state() {
    [[ -r "$STATE_FILE" ]] || return 1
    # shellcheck disable=SC1090
    source "$STATE_FILE"
    return 0
}

get_uid() {
    id -u "$1" 2>/dev/null || err "user '$1' does not exist"
}

iface_exists() {
    ip link show dev "$1" >/dev/null 2>&1
}

iface_is_wireguard() {
    have wg || return 1
    wg show interfaces 2>/dev/null | tr ' ' '\n' | grep -Fxq "$1"
}

detect_single_vpn_iface() {
    local n
    mapfile -t _vpn_candidates < <(
        {
            ip -o link show | awk -F': ' '{print $2}' | grep -E '^(cs-[a-z]+|(tun|wg)[A-Za-z0-9_.-]*)$' || true
            have wg && wg show interfaces 2>/dev/null | tr ' ' '\n' || true
        } | awk 'NF' | sort -u
    )

    n="${#_vpn_candidates[@]}"
    if [[ "$n" -eq 1 ]]; then
        printf '%s\n' "${_vpn_candidates[0]}"
        return 0
    fi
    return 1
}

# Collect directly-connected LAN prefixes from main table, excluding lo and the
# VPN interface. This is more universal than grepping 10./192.168.
collect_lan_prefixes() {
    LAN4=()
    LAN6=()

    if [[ "$ALLOW_LAN" -eq 1 ]]; then
        while IFS= read -r pfx; do
            [[ -n "$pfx" ]] && LAN4+=("$pfx")
        done < <(
            ip -o -4 route show table main scope link 2>/dev/null \
            | awk -v vpn="$VPN_IF" '
                $1 != "default" && $0 !~ / dev lo([[:space:]]|$)/ && $0 !~ (" dev " vpn "([[:space:]]|$)") {print $1}
            ' | sort -u
        )

        while IFS= read -r pfx; do
            [[ -n "$pfx" ]] && LAN6+=("$pfx")
        done < <(
            ip -o -6 route show table main 2>/dev/null \
            | awk -v vpn="$VPN_IF" '
                $1 != "default" &&
                $1 != "::/0" &&
                $1 != "unreachable" &&
                $0 !~ / dev lo([[:space:]]|$)/ &&
                $0 !~ (" dev " vpn "([[:space:]]|$)") &&
                $1 !~ /^fe80::/ {print $1}
            ' | sort -u
        )
    fi
}

backup_resolv_conf() {
    mkdir_state
    if [[ ! -e "$RESOLV_BACKUP" && -e /etc/resolv.conf ]]; then
        cp -L --preserve=mode,ownership,timestamps /etc/resolv.conf "$RESOLV_BACKUP" || err "failed to back up /etc/resolv.conf"
    fi
}

write_vpn_resolv_conf() {
    backup_resolv_conf

    cat > /etc/resolv.conf <<EOF
# managed by vpn-killswitch
nameserver $DNS_V4
nameserver $DNS_V6
options edns0 trust-ad
EOF
    chmod 644 /etc/resolv.conf || err "failed to chmod /etc/resolv.conf"
    RESOLV_CHANGED=1
}

restore_resolv_conf() {
    if [[ -e "$RESOLV_BACKUP" ]]; then
        cp -f "$RESOLV_BACKUP" /etc/resolv.conf
        rm -f "$RESOLV_BACKUP"
    fi
    RESOLV_CHANGED=0
}

###############################################################################
# iptables backend
###############################################################################

ipt()  { iptables  "$@"; }
ip6t() { ip6tables "$@"; }

iptables_add_jump_once() {
    local cmd="$1" chain="$2" jump="$3"
    if ! $cmd -C "$chain" -j "$jump" >/dev/null 2>&1; then
        $cmd -I "$chain" 1 -j "$jump"
    fi
}

iptables_del_jump_if_present() {
    local cmd="$1" chain="$2" jump="$3"
    while $cmd -C "$chain" -j "$jump" >/dev/null 2>&1; do
        $cmd -D "$chain" -j "$jump"
    done
}

iptables_apply_system() {
    collect_lan_prefixes

    # filter chains
    ipt  -N "$IPT_CHAIN_IN"  2>/dev/null || true
    ipt  -N "$IPT_CHAIN_OUT" 2>/dev/null || true
    ip6t -N "$IPT_CHAIN_IN"  2>/dev/null || true
    ip6t -N "$IPT_CHAIN_OUT" 2>/dev/null || true

    ipt  -F "$IPT_CHAIN_IN"
    ipt  -F "$IPT_CHAIN_OUT"
    ip6t -F "$IPT_CHAIN_IN"
    ip6t -F "$IPT_CHAIN_OUT"

    # nat chains
    ipt  -t nat -N "$IPT_CHAIN_NAT" 2>/dev/null || true
    ip6t -t nat -N "$IPT_CHAIN_NAT" 2>/dev/null || true
    ipt  -t nat -F "$IPT_CHAIN_NAT"
    ip6t -t nat -F "$IPT_CHAIN_NAT"

    # INPUT v4/v6
    ipt  -A "$IPT_CHAIN_IN" -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
    ipt  -A "$IPT_CHAIN_IN" -i lo -j ACCEPT
    [[ -n "$VPN_IF" ]] && ipt  -A "$IPT_CHAIN_IN" -i "$VPN_IF" -j ACCEPT
    [[ -n "${ENDPOINT_V4:-}" ]] && ipt -A "$IPT_CHAIN_IN" -s "$ENDPOINT_V4" -j ACCEPT
    for pfx in "${LAN4[@]:-}"; do ipt  -A "$IPT_CHAIN_IN" -s "$pfx" -j ACCEPT; done
    ipt  -A "$IPT_CHAIN_IN" -j REJECT

    ip6t -A "$IPT_CHAIN_IN" -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
    ip6t -A "$IPT_CHAIN_IN" -i lo -j ACCEPT
    [[ -n "$VPN_IF" ]] && ip6t -A "$IPT_CHAIN_IN" -i "$VPN_IF" -j ACCEPT
    [[ -n "${ENDPOINT_V6:-}" ]] && ip6t -A "$IPT_CHAIN_IN" -s "$ENDPOINT_V6" -j ACCEPT
    for pfx in "${LAN6[@]:-}"; do ip6t -A "$IPT_CHAIN_IN" -s "$pfx" -j ACCEPT; done
    ip6t -A "$IPT_CHAIN_IN" -p ipv6-icmp -j ACCEPT
    ip6t -A "$IPT_CHAIN_IN" -j REJECT

    # OUTPUT v4
    ipt  -A "$IPT_CHAIN_OUT" -o lo -j ACCEPT
    [[ -n "$VPN_IF" ]] && ipt  -A "$IPT_CHAIN_OUT" -o "$VPN_IF" -j ACCEPT
    [[ -n "${ENDPOINT_V4:-}" ]] && ipt -A "$IPT_CHAIN_OUT" -d "$ENDPOINT_V4" -j ACCEPT
    [[ -n "$DNS_V4" ]] && ipt  -A "$IPT_CHAIN_OUT" -d "$DNS_V4" -p udp --dport 53 -j ACCEPT
    [[ -n "$DNS_V4" ]] && ipt  -A "$IPT_CHAIN_OUT" -d "$DNS_V4" -p tcp --dport 53 -j ACCEPT
    for pfx in "${LAN4[@]:-}"; do ipt -A "$IPT_CHAIN_OUT" -d "$pfx" -j ACCEPT; done
    ipt  -A "$IPT_CHAIN_OUT" -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
    ipt  -A "$IPT_CHAIN_OUT" -j REJECT

    # OUTPUT v6
    ip6t -A "$IPT_CHAIN_OUT" -o lo -j ACCEPT
    [[ -n "$VPN_IF" ]] && ip6t -A "$IPT_CHAIN_OUT" -o "$VPN_IF" -j ACCEPT
    [[ -n "${ENDPOINT_V6:-}" ]] && ip6t -A "$IPT_CHAIN_OUT" -d "$ENDPOINT_V6" -j ACCEPT
    [[ -n "$DNS_V6" ]] && ip6t -A "$IPT_CHAIN_OUT" -d "$DNS_V6" -p udp --dport 53 -j ACCEPT
    [[ -n "$DNS_V6" ]] && ip6t -A "$IPT_CHAIN_OUT" -d "$DNS_V6" -p tcp --dport 53 -j ACCEPT
    for pfx in "${LAN6[@]:-}"; do ip6t -A "$IPT_CHAIN_OUT" -d "$pfx" -j ACCEPT; done
    ip6t -A "$IPT_CHAIN_OUT" -p ipv6-icmp -j ACCEPT
    ip6t -A "$IPT_CHAIN_OUT" -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
    ip6t -A "$IPT_CHAIN_OUT" -j REJECT

    # DNS DNAT for hardcoded resolvers; resolv.conf rewrite handles 127.0.0.53
    [[ -n "$DNS_V4" ]] && ipt  -t nat -A "$IPT_CHAIN_NAT" ! -o lo -p udp --dport 53 -j DNAT --to-destination "$DNS_V4"
    [[ -n "$DNS_V4" ]] && ipt  -t nat -A "$IPT_CHAIN_NAT" ! -o lo -p tcp --dport 53 -j DNAT --to-destination "$DNS_V4"
    [[ -n "$DNS_V6" ]] && ip6t -t nat -A "$IPT_CHAIN_NAT" ! -o lo -p udp --dport 53 -j DNAT --to-destination "$DNS_V6"
    [[ -n "$DNS_V6" ]] && ip6t -t nat -A "$IPT_CHAIN_NAT" ! -o lo -p tcp --dport 53 -j DNAT --to-destination "$DNS_V6"

    # Install jumps
    iptables_add_jump_once ipt  INPUT  "$IPT_CHAIN_IN"
    iptables_add_jump_once ipt  OUTPUT "$IPT_CHAIN_OUT"
    iptables_add_jump_once ip6t INPUT  "$IPT_CHAIN_IN"
    iptables_add_jump_once ip6t OUTPUT "$IPT_CHAIN_OUT"
    iptables_add_jump_once "iptables -t nat"  OUTPUT "$IPT_CHAIN_NAT"
    iptables_add_jump_once "ip6tables -t nat" OUTPUT "$IPT_CHAIN_NAT"
}

iptables_apply_user() {
    local uid
    uid="$(get_uid "$KS_USER")"

    # OUTPUT user-only filter chains
    ipt  -N "$IPT_CHAIN_OUT_USER" 2>/dev/null || true
    ip6t -N "$IPT_CHAIN_OUT_USER" 2>/dev/null || true
    ipt  -F "$IPT_CHAIN_OUT_USER"
    ip6t -F "$IPT_CHAIN_OUT_USER"

    # nat chains
    ipt  -t nat -N "$IPT_CHAIN_NAT" 2>/dev/null || true
    ip6t -t nat -N "$IPT_CHAIN_NAT" 2>/dev/null || true
    ipt  -t nat -F "$IPT_CHAIN_NAT"
    ip6t -t nat -F "$IPT_CHAIN_NAT"

    collect_lan_prefixes

    ipt  -A "$IPT_CHAIN_OUT_USER" -o lo -j ACCEPT
    [[ -n "$VPN_IF" ]] && ipt  -A "$IPT_CHAIN_OUT_USER" -o "$VPN_IF" -j ACCEPT
    for pfx in "${LAN4[@]:-}"; do ipt -A "$IPT_CHAIN_OUT_USER" -d "$pfx" -j ACCEPT; done
    [[ -n "${ENDPOINT_V4:-}" ]] && ipt -A "$IPT_CHAIN_OUT_USER" -d "$ENDPOINT_V4" -j ACCEPT
    [[ -n "$DNS_V4" ]] && ipt -A "$IPT_CHAIN_OUT_USER" -d "$DNS_V4" -p udp --dport 53 -j ACCEPT
    [[ -n "$DNS_V4" ]] && ipt -A "$IPT_CHAIN_OUT_USER" -d "$DNS_V4" -p tcp --dport 53 -j ACCEPT
    ipt  -A "$IPT_CHAIN_OUT_USER" -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
    ipt  -A "$IPT_CHAIN_OUT_USER" -j REJECT

    ip6t -A "$IPT_CHAIN_OUT_USER" -o lo -j ACCEPT
    [[ -n "$VPN_IF" ]] && ip6t -A "$IPT_CHAIN_OUT_USER" -o "$VPN_IF" -j ACCEPT
    for pfx in "${LAN6[@]:-}"; do ip6t -A "$IPT_CHAIN_OUT_USER" -d "$pfx" -j ACCEPT; done
    [[ -n "${ENDPOINT_V6:-}" ]] && ip6t -A "$IPT_CHAIN_OUT_USER" -d "$ENDPOINT_V6" -j ACCEPT
    [[ -n "$DNS_V6" ]] && ip6t -A "$IPT_CHAIN_OUT_USER" -d "$DNS_V6" -p udp --dport 53 -j ACCEPT
    [[ -n "$DNS_V6" ]] && ip6t -A "$IPT_CHAIN_OUT_USER" -d "$DNS_V6" -p tcp --dport 53 -j ACCEPT
    ip6t -A "$IPT_CHAIN_OUT_USER" -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
    ip6t -A "$IPT_CHAIN_OUT_USER" -j REJECT

    [[ -n "$DNS_V4" ]] && ipt  -t nat -A "$IPT_CHAIN_NAT" -m owner --uid-owner "$uid" ! -o lo -p udp --dport 53 -j DNAT --to-destination "$DNS_V4"
    [[ -n "$DNS_V4" ]] && ipt  -t nat -A "$IPT_CHAIN_NAT" -m owner --uid-owner "$uid" ! -o lo -p tcp --dport 53 -j DNAT --to-destination "$DNS_V4"
    [[ -n "$DNS_V6" ]] && ip6t -t nat -A "$IPT_CHAIN_NAT" -m owner --uid-owner "$uid" ! -o lo -p udp --dport 53 -j DNAT --to-destination "$DNS_V6"
    [[ -n "$DNS_V6" ]] && ip6t -t nat -A "$IPT_CHAIN_NAT" -m owner --uid-owner "$uid" ! -o lo -p tcp --dport 53 -j DNAT --to-destination "$DNS_V6"

    if ! ipt -C OUTPUT -m owner --uid-owner "$uid" -j "$IPT_CHAIN_OUT_USER" >/dev/null 2>&1; then
        ipt -I OUTPUT 1 -m owner --uid-owner "$uid" -j "$IPT_CHAIN_OUT_USER"
    fi
    if ! ip6t -C OUTPUT -m owner --uid-owner "$uid" -j "$IPT_CHAIN_OUT_USER" >/dev/null 2>&1; then
        ip6t -I OUTPUT 1 -m owner --uid-owner "$uid" -j "$IPT_CHAIN_OUT_USER"
    fi

    iptables_add_jump_once "iptables -t nat"  OUTPUT "$IPT_CHAIN_NAT"
    iptables_add_jump_once "ip6tables -t nat" OUTPUT "$IPT_CHAIN_NAT"

    # Route only that user's traffic into the tunnel
    [[ -n "$VPN_IF" ]] || err "MODE=user needs a VPN interface"

    ip rule add pref "$RULE_PREF" uidrange "$uid-$uid" lookup "$ROUTE_TABLE_ID" 2>/dev/null || true
    ip route replace table "$ROUTE_TABLE_ID" default dev "$VPN_IF"

    # Add IPv6 rule/table if the tunnel has global IPv6 or DNS_V6 is configured
    ip -6 rule add pref "$RULE_PREF" uidrange "$uid-$uid" lookup "$ROUTE_TABLE_ID" 2>/dev/null || true
    ip -6 route replace table "$ROUTE_TABLE_ID" default dev "$VPN_IF" 2>/dev/null || true
    ip route flush cache 2>/dev/null || true
}

iptables_remove_all() {
    if load_state; then
        if [[ "$MODE" == "user" ]]; then
            local uid
            uid="$(get_uid "$KS_USER")" || true

            while ipt -C OUTPUT -m owner --uid-owner "$uid" -j "$IPT_CHAIN_OUT_USER" >/dev/null 2>&1; do
                ipt -D OUTPUT -m owner --uid-owner "$uid" -j "$IPT_CHAIN_OUT_USER"
            done
            while ip6t -C OUTPUT -m owner --uid-owner "$uid" -j "$IPT_CHAIN_OUT_USER" >/dev/null 2>&1; do
                ip6t -D OUTPUT -m owner --uid-owner "$uid" -j "$IPT_CHAIN_OUT_USER"
            done

            ip rule del pref "$RULE_PREF" uidrange "$uid-$uid" lookup "$ROUTE_TABLE_ID" 2>/dev/null || true
            ip -6 rule del pref "$RULE_PREF" uidrange "$uid-$uid" lookup "$ROUTE_TABLE_ID" 2>/dev/null || true
            ip route flush table "$ROUTE_TABLE_ID" 2>/dev/null || true
            ip -6 route flush table "$ROUTE_TABLE_ID" 2>/dev/null || true
        else
            iptables_del_jump_if_present ipt  INPUT  "$IPT_CHAIN_IN"
            iptables_del_jump_if_present ipt  OUTPUT "$IPT_CHAIN_OUT"
            iptables_del_jump_if_present ip6t INPUT  "$IPT_CHAIN_IN"
            iptables_del_jump_if_present ip6t OUTPUT "$IPT_CHAIN_OUT"
        fi

        iptables_del_jump_if_present "iptables -t nat"  OUTPUT "$IPT_CHAIN_NAT"
        iptables_del_jump_if_present "ip6tables -t nat" OUTPUT "$IPT_CHAIN_NAT"
    fi

    # Flush/remove our chains if they exist
    ipt  -F "$IPT_CHAIN_IN"       2>/dev/null || true
    ipt  -F "$IPT_CHAIN_OUT"      2>/dev/null || true
    ipt  -F "$IPT_CHAIN_OUT_USER" 2>/dev/null || true
    ipt  -t nat -F "$IPT_CHAIN_NAT" 2>/dev/null || true
    ipt  -X "$IPT_CHAIN_IN"       2>/dev/null || true
    ipt  -X "$IPT_CHAIN_OUT"      2>/dev/null || true
    ipt  -X "$IPT_CHAIN_OUT_USER" 2>/dev/null || true
    ipt  -t nat -X "$IPT_CHAIN_NAT" 2>/dev/null || true

    ip6t -F "$IPT_CHAIN_IN"       2>/dev/null || true
    ip6t -F "$IPT_CHAIN_OUT"      2>/dev/null || true
    ip6t -F "$IPT_CHAIN_OUT_USER" 2>/dev/null || true
    ip6t -t nat -F "$IPT_CHAIN_NAT" 2>/dev/null || true
    ip6t -X "$IPT_CHAIN_IN"       2>/dev/null || true
    ip6t -X "$IPT_CHAIN_OUT"      2>/dev/null || true
    ip6t -X "$IPT_CHAIN_OUT_USER" 2>/dev/null || true
    ip6t -t nat -X "$IPT_CHAIN_NAT" 2>/dev/null || true
}

###############################################################################
# nft backend
###############################################################################
        
nft_delete_tables_if_exist() {
    nft delete table inet "$NFT_TABLE_INET" 2>/dev/null || true
    nft delete table ip "$NFT_TABLE_IP_NAT" 2>/dev/null || true
    nft delete table ip6 "$NFT_TABLE_IP6_NAT" 2>/dev/null || true
}

nft_apply_system() {
    collect_lan_prefixes
    nft_delete_tables_if_exist

    local nftfile
    nftfile="$(mktemp)" || err "failed to create temporary nft file"
    NFT_TMPFILE="$nftfile"
    {
        echo "table inet $NFT_TABLE_INET {"

        if ((${#LAN4[@]})); then
            printf '  set lan4 { type ipv4_addr; flags interval; elements = { '
            printf '%s, ' "${LAN4[@]}" | sed 's/, $//'
            echo ' } }'
        fi
        if ((${#LAN6[@]})); then
            printf '  set lan6 { type ipv6_addr; flags interval; elements = { '
            printf '%s, ' "${LAN6[@]}" | sed 's/, $//'
            echo ' } }'
        fi
        if [[ -n "${ENDPOINT_V4:-}" ]]; then
            echo "  set endpoint4 { type ipv4_addr; elements = { $ENDPOINT_V4 } }"
        fi
        if [[ -n "${ENDPOINT_V6:-}" ]]; then
            echo "  set endpoint6 { type ipv6_addr; elements = { $ENDPOINT_V6 } }"
        fi

        cat <<EOF
  chain input {
    type filter hook input priority 0; policy accept;
    ct state established,related accept
    iifname "lo" accept
EOF
        [[ -n "$VPN_IF" ]] && echo "    iifname \"$VPN_IF\" accept"
        [[ -n "${ENDPOINT_V4:-}" ]] && echo "    ip saddr @endpoint4 accept"
        [[ -n "${ENDPOINT_V6:-}" ]] && echo "    ip6 saddr @endpoint6 accept"
        ((${#LAN4[@]})) && echo "    ip saddr @lan4 accept"
        ((${#LAN6[@]})) && echo "    ip6 saddr @lan6 accept"
        echo "    meta l4proto ipv6-icmp accept"
        echo "    reject"
        echo "  }"

        cat <<EOF
  chain output {
    type filter hook output priority 0; policy accept;
    oifname "lo" accept
EOF
        [[ -n "$VPN_IF" ]] && echo "    oifname \"$VPN_IF\" accept"
        [[ -n "${ENDPOINT_V4:-}" ]] && echo "    ip daddr @endpoint4 accept"
        [[ -n "${ENDPOINT_V6:-}" ]] && echo "    ip6 daddr @endpoint6 accept"
        [[ -n "$DNS_V4" ]] && echo "    ip daddr $DNS_V4 udp dport 53 accept"
        [[ -n "$DNS_V4" ]] && echo "    ip daddr $DNS_V4 tcp dport 53 accept"
        [[ -n "$DNS_V6" ]] && echo "    ip6 daddr $DNS_V6 udp dport 53 accept"
        [[ -n "$DNS_V6" ]] && echo "    ip6 daddr $DNS_V6 tcp dport 53 accept"
        ((${#LAN4[@]})) && echo "    ip daddr @lan4 accept"
        ((${#LAN6[@]})) && echo "    ip6 daddr @lan6 accept"
        echo "    meta l4proto ipv6-icmp accept"
        echo "    ct state established,related accept"
        echo "    reject"
        echo "  }"
        echo "}"

        echo "table ip $NFT_TABLE_IP_NAT {"
        echo "  chain output { type nat hook output priority -100; policy accept;"
        [[ -n "$DNS_V4" ]] && echo "    oifname != \"lo\" udp dport 53 dnat to $DNS_V4"
        [[ -n "$DNS_V4" ]] && echo "    oifname != \"lo\" tcp dport 53 dnat to $DNS_V4"
        echo "  }"
        echo "}"

        echo "table ip6 $NFT_TABLE_IP6_NAT {"
        echo "  chain output { type nat hook output priority -100; policy accept;"
        [[ -n "$DNS_V6" ]] && echo "    oifname != \"lo\" udp dport 53 dnat to $DNS_V6"
        [[ -n "$DNS_V6" ]] && echo "    oifname != \"lo\" tcp dport 53 dnat to $DNS_V6"
        echo "  }"
        echo "}"
    } > "$nftfile"

    nft -f "$nftfile"
    cleanup_temp_files
}

nft_apply_user() {
    local uid
    uid="$(get_uid "$KS_USER")"
    collect_lan_prefixes
    nft_delete_tables_if_exist

    local nftfile
    nftfile="$(mktemp)" || err "failed to create temporary nft file"
    NFT_TMPFILE="$nftfile"
    {
        echo "table inet $NFT_TABLE_INET {"

        if ((${#LAN4[@]})); then
            printf '  set lan4 { type ipv4_addr; flags interval; elements = { '
            printf '%s, ' "${LAN4[@]}" | sed 's/, $//'
            echo ' } }'
        fi
        if ((${#LAN6[@]})); then
            printf '  set lan6 { type ipv6_addr; flags interval; elements = { '
            printf '%s, ' "${LAN6[@]}" | sed 's/, $//'
            echo ' } }'
        fi
        if [[ -n "${ENDPOINT_V4:-}" ]]; then
            echo "  set endpoint4 { type ipv4_addr; elements = { $ENDPOINT_V4 } }"
        fi
        if [[ -n "${ENDPOINT_V6:-}" ]]; then
            echo "  set endpoint6 { type ipv6_addr; elements = { $ENDPOINT_V6 } }"
        fi

        cat <<EOF
  chain output {
    type filter hook output priority 0; policy accept;
    meta skuid != $uid accept
    oifname "lo" accept
EOF
        [[ -n "$VPN_IF" ]] && echo "    oifname \"$VPN_IF\" accept"
        [[ -n "${ENDPOINT_V4:-}" ]] && echo "    ip daddr @endpoint4 accept"
        [[ -n "${ENDPOINT_V6:-}" ]] && echo "    ip6 daddr @endpoint6 accept"
        [[ -n "$DNS_V4" ]] && echo "    ip daddr $DNS_V4 udp dport 53 accept"
        [[ -n "$DNS_V4" ]] && echo "    ip daddr $DNS_V4 tcp dport 53 accept"
        [[ -n "$DNS_V6" ]] && echo "    ip6 daddr $DNS_V6 udp dport 53 accept"
        [[ -n "$DNS_V6" ]] && echo "    ip6 daddr $DNS_V6 tcp dport 53 accept"
        ((${#LAN4[@]})) && echo "    ip daddr @lan4 accept"
        ((${#LAN6[@]})) && echo "    ip6 daddr @lan6 accept"
        echo "    ct state established,related accept"
        echo "    reject"
        echo "  }"
        echo "}"
        echo "table ip $NFT_TABLE_IP_NAT {"
        echo "  chain output { type nat hook output priority -100; policy accept;"
        [[ -n "$DNS_V4" ]] && echo "    meta skuid $uid oifname != \"lo\" udp dport 53 dnat to $DNS_V4"
        [[ -n "$DNS_V4" ]] && echo "    meta skuid $uid oifname != \"lo\" tcp dport 53 dnat to $DNS_V4"
        echo "  }"
        echo "}"
        echo "table ip6 $NFT_TABLE_IP6_NAT {"
        echo "  chain output { type nat hook output priority -100; policy accept;"
        [[ -n "$DNS_V6" ]] && echo "    meta skuid $uid oifname != \"lo\" udp dport 53 dnat to $DNS_V6"
        [[ -n "$DNS_V6" ]] && echo "    meta skuid $uid oifname != \"lo\" tcp dport 53 dnat to $DNS_V6"
        echo "  }"
        echo "}"
    } > "$nftfile"

    nft -f "$nftfile"
    cleanup_temp_files

    [[ -n "$VPN_IF" ]] || err "MODE=user needs a VPN interface"

    ip rule add pref "$RULE_PREF" uidrange "$uid-$uid" lookup "$ROUTE_TABLE_ID" 2>/dev/null || true
    ip route replace table "$ROUTE_TABLE_ID" default dev "$VPN_IF"
    ip -6 rule add pref "$RULE_PREF" uidrange "$uid-$uid" lookup "$ROUTE_TABLE_ID" 2>/dev/null || true
    ip -6 route replace table "$ROUTE_TABLE_ID" default dev "$VPN_IF" 2>/dev/null || true
    ip route flush cache 2>/dev/null || true
}

nft_remove_all() {
    if load_state && [[ "$MODE" == "user" ]]; then
        local uid
        uid="$(get_uid "$KS_USER")" || true
        ip rule del pref "$RULE_PREF" uidrange "$uid-$uid" lookup "$ROUTE_TABLE_ID" 2>/dev/null || true
        ip -6 rule del pref "$RULE_PREF" uidrange "$uid-$uid" lookup "$ROUTE_TABLE_ID" 2>/dev/null || true
        ip route flush table "$ROUTE_TABLE_ID" 2>/dev/null || true
        ip -6 route flush table "$ROUTE_TABLE_ID" 2>/dev/null || true
    fi

    nft delete table inet "$NFT_TABLE_INET" 2>/dev/null || true
    nft delete table ip "$NFT_TABLE_IP_NAT" 2>/dev/null || true
    nft delete table ip6 "$NFT_TABLE_IP6_NAT" 2>/dev/null || true
}

###############################################################################
# Main
###############################################################################

enable_killswitch() {
    ENABLE_IN_PROGRESS=1
    choose_backend
    check_backend_capabilities
    resolve_vpn_endpoints
    
    [[ -n "${VPN_IF:-}" ]] || err "could not determine VPN interface"

    iface_exists "$VPN_IF" || err "interface '$VPN_IF' does not exist"
      
    if iface_is_wireguard "$VPN_IF"; then
        wait_for_wireguard_ready "$VPN_IF"
        resolve_vpn_endpoints
    fi
    
    # Standalone OpenVPN mode needs a known endpoint. In --up/--down mode OpenVPN
    # provides trusted_ip / trusted_ip6, but in standalone mode it usually does not.
    if iface_is_openvpn_tun "$VPN_IF" \
       && [[ -z "${script_type:-}" ]] \
       && [[ -z "${ENDPOINT_V4:-}" && -z "${ENDPOINT_V6:-}" ]]; then
        err "standalone OpenVPN mode requires VPN_HOST or VPN_ENDPOINT_V4/VPN_ENDPOINT_V6 to be set; otherwise the server endpoint cannot be safely exempted"
    fi

    # In MODE=system we really want at least the tunnel interface.
    if [[ "$MODE" == "system" && -z "${VPN_IF:-}" ]]; then
        err "could not determine VPN interface; set VPN_IF_OVERRIDE or use OpenVPN/NM hook mode"
    fi

    write_vpn_resolv_conf
    save_state

    if [[ "$FW_BACKEND" == "nft" ]]; then
        if [[ "$MODE" == "user" ]]; then
            nft_apply_user
        else
            nft_apply_system
        fi
    else
        if [[ "$MODE" == "user" ]]; then
            iptables_apply_user
        else
            iptables_apply_system
        fi
    fi

    ENABLE_IN_PROGRESS=0
    log "Killswitch enabled (${MODE}, backend=${FW_BACKEND}, vpn_if=${VPN_IF:-unknown})"
}

disable_killswitch() {
    if ! load_state; then
        log "No saved state found; attempting best-effort cleanup."
        choose_backend
    else
        FW_BACKEND="${FW_BACKEND:-$BACKEND}"
    fi

    if [[ "${FW_BACKEND:-}" == "nft" ]]; then
        nft_remove_all
    else
        iptables_remove_all
    fi

    restore_resolv_conf
    rm -f "$STATE_FILE"
    log "Killswitch disabled"
}

status_killswitch() {
    if load_state; then
        cat <<EOF
enabled=yes
backend=$FW_BACKEND
mode=$MODE
user=${KS_USER:-}
allow_lan=$ALLOW_LAN
vpn_if=${VPN_IF:-}
dns_v4=${DNS_V4:-}
dns_v6=${DNS_V6:-}
endpoint_v4=${ENDPOINT_V4:-}
endpoint_v6=${ENDPOINT_V6:-}
EOF
    else
        echo "enabled=no"
    fi
}

need_root
[[ "$(uname -s)" == "Linux" ]] || err "this script is for Linux"

trap 'rc=$?; cleanup_temp_files; if [[ "${ENABLE_IN_PROGRESS:-0}" -eq 1 && $rc -ne 0 ]]; then rollback_enable; fi; exit $rc' EXIT
trap 'rc=$?; cleanup_temp_files; if [[ "${ENABLE_IN_PROGRESS:-0}" -eq 1 && $rc -ne 0 ]]; then rollback_enable; fi; exit $rc' ERR
trap 'cleanup_temp_files' INT TERM

detect_action_and_iface "$@"

case "$ACTION" in
    enable)  enable_killswitch ;;
    disable) disable_killswitch ;;
    status)  status_killswitch ;;
    *) err "unknown action" ;;
esac