#!/bin/bash
set -euo pipefail

# ----------------------------
# Config
# ----------------------------
MIHOMO_UID="mihomo"
REDIR_PORT="7892"     # mihomo redir-port (NAT REDIRECT for TCP)
TPROXY_PORT="7893"    # mihomo tproxy-port (TPROXY for TCP/UDP)
FW_MARK="0x1"
ROUTE_TABLE="100"

# Interfaces to EXCLUDE completely from interception
EXCLUDE_IFACES=("tun0" "wg0")

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

# Функция для удаления правил по комментарию (чтобы не тереть чужое)
# Используем trick: добавляем комментарий ко всем правилам, чтобы потом их найти и удалить
cleanup_rules() {
  local table="$1"
  local chain="$2"
  local comment="$3"
  
  # List rules with line numbers, grep by comment, sort reverse order (to delete safely), delete
  iptables -t "$table" -nL "$chain" --line-numbers 2>/dev/null | \
  grep "$comment" | \
  sort -r | \
  awk '{print $1}' | \
  while read -r line; do
    echo "Deleting rule $line from $table/$chain..."
    iptables -t "$table" -D "$chain" "$line"
  done
}

ensure_ip_rule() {
  # Clean up ip rules first
  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}
  # Route table entry, forced
  ip route replace local 0.0.0.0/0 dev lo table ${ROUTE_TABLE}
}

# ----------------------------
# CLEANUP PHASE
# ----------------------------
echo "Cleaning up old rules..."

# 1. Clean MANGLE references
cleanup_rules mangle OUTPUT "MIHOMO-MARK"
cleanup_rules mangle OUTPUT "MIHOMO-EXCLUDE"
cleanup_rules mangle PREROUTING "MIHOMO-JUMP" 
cleanup_rules mangle PREROUTING "MIHOMO-EXCLUDE"

# 2. Clean NAT references
cleanup_rules nat OUTPUT "MIHOMO-JUMP"
cleanup_rules nat OUTPUT "MIHOMO-EXCLUDE"
cleanup_rules nat PREROUTING "MIHOMO-REDIRECT"

# 3. Flush and Delete Chains
# Now that references are gone, we can safely kill the chains
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

# ----------------------------
# NAT (REDIRECT) part (TCP only)
# ----------------------------
echo "Applying NAT rules..."

ipt -t nat -N MIHOMO_REDIR

# Exclude private networks
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

# Everything else TCP -> REDIRECT
ipt -t nat -A MIHOMO_REDIR -p tcp -j REDIRECT --to-ports "${REDIR_PORT}"

# --- APPLY TO OUTPUT (Local Traffic) ---
# Exclude traffic going via tunnels (tun0/wg0)
for IFACE in "${EXCLUDE_IFACES[@]}"; do
  ipt -t nat -A OUTPUT -o "${IFACE}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN
done

# Exclude mihomo user
ipt -t nat -A OUTPUT -m owner --uid-owner "${MIHOMO_UID}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN

# Redirect remaining TCP
ipt -t nat -A OUTPUT -p tcp -m comment --comment "MIHOMO-JUMP" -j MIHOMO_REDIR

# --- APPLY TO INGRESS (wt0) ---
# Exclude wt0 -> tun0/wg0 logic handled by routing mostly, but let's be safe if needed.
# For now, strictly redirect TCP incoming on wt0
ipt -t nat -A PREROUTING -i wt0 -p tcp -m comment --comment "MIHOMO-REDIRECT" -j REDIRECT --to-port "${REDIR_PORT}"


# ----------------------------
# MANGLE (TPROXY) part (TCP+UDP)
# ----------------------------
echo "Applying TPROXY rules..."

ensure_ip_rule

ipt -t mangle -N MIHOMO_TPROXY

# Exclude private networks
ipt -t mangle -A MIHOMO_TPROXY -d 192.168.0.0/16 -j RETURN
ipt -t mangle -A MIHOMO_TPROXY -d 10.0.0.0/8 -j RETURN
ipt -t mangle -A MIHOMO_TPROXY -d 172.16.0.0/12 -j RETURN
ipt -t mangle -A MIHOMO_TPROXY -d 127.0.0.0/8 -j RETURN

# TPROXY Target
ipt -t mangle -A MIHOMO_TPROXY -p tcp -j TPROXY --on-port "${TPROXY_PORT}" --tproxy-mark "${FW_MARK}/${FW_MARK}"
ipt -t mangle -A MIHOMO_TPROXY -p udp -j TPROXY --on-port "${TPROXY_PORT}" --tproxy-mark "${FW_MARK}/${FW_MARK}"

# --- APPLY TO OUTPUT (Local Traffic) ---
# Exclude tunnels
for IFACE in "${EXCLUDE_IFACES[@]}"; do
  ipt -t mangle -A OUTPUT -o "${IFACE}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN
done

# Exclude mihomo user
ipt -t mangle -A OUTPUT -m owner --uid-owner "${MIHOMO_UID}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN

# Mark remaining traffic for Policy Routing
ipt -t mangle -A OUTPUT -p tcp -m comment --comment "MIHOMO-MARK" -j MARK --set-mark "${FW_MARK}"
ipt -t mangle -A OUTPUT -p udp -m comment --comment "MIHOMO-MARK" -j MARK --set-mark "${FW_MARK}"

# --- APPLY TO INGRESS (wt0) ---
# Exclude tunnels (ingress logic)
for IFACE in "${EXCLUDE_IFACES[@]}"; do
  ipt -t mangle -A PREROUTING -i "${IFACE}" -m comment --comment "MIHOMO-EXCLUDE" -j RETURN
done

# Jump wt0 traffic to TPROXY chain
ipt -t mangle -A PREROUTING -i wt0 -m comment --comment "MIHOMO-JUMP" -j MIHOMO_TPROXY

echo "Done."