#!/bin/bash
set -euo pipefail

# ----------------------------
# Config
# ----------------------------
MIHOMO_UID="mihomo"
REDIR_PORT="7892"     # TCP Redirect
TPROXY_PORT="7893"    # UDP/TCP TProxy
FW_MARK="0x1"
ROUTE_TABLE="100"

EXCLUDE_IFACES=("tun0")  
INCLUDE_IFACES=("wt0" "eth1" "eth2")

# ----------------------------
# Helpers
# ----------------------------
ipt() { iptables "$@"; }

del_loop() {
  local table=$1
  local chain=$2
  shift 2
  local rule_args="$@"

  while iptables -t "$table" -C "$chain" $rule_args 2>/dev/null; do
    echo "Deleting from $table/$chain: $rule_args"
    iptables -t "$table" -D "$chain" $rule_args
  done
}

ensure_ip_rule() {
  while ip rule list | grep -q "fwmark ${FW_MARK} lookup ${ROUTE_TABLE}"; do
    ip rule del fwmark ${FW_MARK} lookup ${ROUTE_TABLE} || true
  done
  ip rule add fwmark ${FW_MARK} lookup ${ROUTE_TABLE}
  ip route replace local 0.0.0.0/0 dev lo table ${ROUTE_TABLE}
}

# ----------------------------
# CLEANUP PHASE
# ----------------------------
echo "--- Cleaning up old rules (Robust Mode) ---"

del_loop nat OUTPUT -p tcp -m comment --comment "MIHOMO-JUMP" -j MIHOMO_REDIR
del_loop nat PREROUTING -i wt0 -p tcp -m comment --comment "MIHOMO-REDIRECT" -j REDIRECT --to-port "${REDIR_PORT}"
del_loop mangle PREROUTING -i wt0 -m comment --comment "MIHOMO-JUMP" -j MIHOMO_TPROXY

del_loop mangle OUTPUT -p tcp -m comment --comment "MIHOMO-MARK" -j MARK --set-mark "${FW_MARK}"
del_loop mangle OUTPUT -p udp -m comment --comment "MIHOMO-MARK" -j MARK --set-mark "${FW_MARK}"

for IFACE in "${EXCLUDE_IFACES[@]}"; do
  del_loop mangle OUTPUT -o "${IFACE}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN
  del_loop mangle PREROUTING -i "${IFACE}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN
  del_loop nat OUTPUT -o "${IFACE}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN
done

del_loop mangle OUTPUT -m owner --uid-owner "${MIHOMO_UID}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN
del_loop nat OUTPUT -m owner --uid-owner "${MIHOMO_UID}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN

ipt -t mangle -F MIHOMO_TPROXY 2>/dev/null || true
ipt -t mangle -X MIHOMO_TPROXY 2>/dev/null || true

ipt -t nat -F MIHOMO_REDIR 2>/dev/null || true
ipt -t nat -X MIHOMO_REDIR 2>/dev/null || true

echo "--- Cleanup finished. Applying new rules ---"

# ----------------------------
# NAT (REDIRECT) - TCP
# ----------------------------
ipt -t nat -N MIHOMO_REDIR

# Exclusions for gateway's own traffic
ipt -t nat -A MIHOMO_REDIR -d 192.168.0.0/16 -j RETURN
ipt -t nat -A MIHOMO_REDIR -d 10.0.0.0/8 -j RETURN
ipt -t nat -A MIHOMO_REDIR -d 172.16.0.0/12 -j RETURN
ipt -t nat -A MIHOMO_REDIR -d 127.0.0.0/8 -j RETURN
ipt -t nat -A MIHOMO_REDIR -p tcp -j REDIRECT --to-ports "${REDIR_PORT}"

# Apply to OUTPUT (Local gateway traffic)
for IFACE in "${EXCLUDE_IFACES[@]}"; do
  ipt -t nat -A OUTPUT -o "${IFACE}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN
done
ipt -t nat -A OUTPUT -m owner --uid-owner "${MIHOMO_UID}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN
ipt -t nat -A OUTPUT -p tcp -m comment --comment "MIHOMO-JUMP" -j MIHOMO_REDIR

# Apply to PREROUTING (wt0 Ingress) - Force Redir for NetBird (skips exclusions by design)
for IFACE in "${INCLUDE_IFACES[@]}"; do
  if [ "$IFACE" = "wt0" ]; then
    # wt0 (Netbird) пропускает исключения локальных подсетей по твоему дизайну
    ipt -t nat -A PREROUTING -i "$IFACE" -p tcp -m comment --comment "MIHOMO-REDIRECT" -j REDIRECT --to-port "${REDIR_PORT}"
  else
    # LAN-трафик (eth1, eth2) должен прыгать в цепочку MIHOMO_REDIR для проверки исключений (192.168.x.x и т.д.)
    ipt -t nat -A PREROUTING -i "$IFACE" -p tcp -m comment --comment "MIHOMO-JUMP-$IFACE" -j MIHOMO_REDIR
  fi
done

# ----------------------------
# MANGLE (TPROXY) - UDP
# ----------------------------
ensure_ip_rule
ipt -t mangle -N MIHOMO_TPROXY

# Local exclusions: apply ONLY if traffic is NOT coming from NetBird (wt0)
ipt -t mangle -A MIHOMO_TPROXY ! -i wt0 -d 192.168.0.0/16 -j RETURN
ipt -t mangle -A MIHOMO_TPROXY ! -i wt0 -d 10.0.0.0/8 -j RETURN
ipt -t mangle -A MIHOMO_TPROXY ! -i wt0 -d 172.16.0.0/12 -j RETURN
ipt -t mangle -A MIHOMO_TPROXY ! -i wt0 -d 127.0.0.0/8 -j RETURN

# TProxy Targets (UDP only, TCP is handled by REDIRECT)
ipt -t mangle -A MIHOMO_TPROXY -p udp -j TPROXY --on-port "${TPROXY_PORT}" --tproxy-mark "${FW_MARK}/${FW_MARK}"

# Apply to OUTPUT (Local gateway traffic)
for IFACE in "${EXCLUDE_IFACES[@]}"; do
  ipt -t mangle -A OUTPUT -o "${IFACE}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN
done
ipt -t mangle -A OUTPUT -m owner --uid-owner "${MIHOMO_UID}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN

# Mark local UDP packets
ipt -t mangle -A OUTPUT -p udp -m comment --comment "MIHOMO-MARK" -j MARK --set-mark "${FW_MARK}"

# Apply to PREROUTING (wt0 Ingress)
for IFACE in "${EXCLUDE_IFACES[@]}"; do
  ipt -t mangle -A PREROUTING -i "${IFACE}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN
done

for IFACE in "${INCLUDE_IFACES[@]}"; do
  ipt -t mangle -A PREROUTING -i "$IFACE" -m comment --comment "MIHOMO-JUMP-$IFACE" -j MIHOMO_TPROXY
done

echo "Done. Suboptimal hypervisor constraints bypassed successfully."