#!/bin/bash
# sub2api 指纹防泄露 iptables 规则
# 确保只有 Node.js TLS Proxy 能直连上游 HTTPS，
# sub2api Go 进程即使有 bug 也无法绕过。
#
# 用法:
#   sudo bash setup-firewall.sh [apply|remove|status]
#
# 前置条件:
#   - Node.js proxy 以专用用户 "nodeproxy" 运行
#   - 创建用户: sudo useradd -r -s /usr/sbin/nologin nodeproxy

set -euo pipefail

NODE_PROXY_USER="${MG_NODE_PROXY_USER:-nodeproxy}"
CHAIN_NAME="MG_FINGERPRINT"

log() { echo "[$(date '+%H:%M:%S')] $*"; }

apply_rules() {
    log "Applying fingerprint firewall rules..."

    # 验证用户存在
    if ! id "$NODE_PROXY_USER" &>/dev/null; then
        log "ERROR: User '$NODE_PROXY_USER' does not exist."
        log "Create it: sudo useradd -r -s /usr/sbin/nologin $NODE_PROXY_USER"
        exit 1
    fi

    # 创建自定义链（幂等）
    iptables -N "$CHAIN_NAME" 2>/dev/null || iptables -F "$CHAIN_NAME"

    # === Rule 1: QUIC 阻断 — 丢弃所有出站 UDP 443/4433 ===
    iptables -A "$CHAIN_NAME" -p udp --dport 443  -j DROP \
        -m comment --comment "MG: block QUIC/HTTP3 UDP 443"
    iptables -A "$CHAIN_NAME" -p udp --dport 4433 -j DROP \
        -m comment --comment "MG: block QUIC alt UDP 4433"

    # === Rule 2: 允许 Node.js proxy 出站 TCP 443 ===
    iptables -A "$CHAIN_NAME" -p tcp --dport 443 \
        -m owner --uid-owner "$NODE_PROXY_USER" -j ACCEPT \
        -m comment --comment "MG: allow nodeproxy TCP 443"

    # === Rule 3: 阻止其他进程直连 TCP 443 ===
    iptables -A "$CHAIN_NAME" -p tcp --dport 443 -j REJECT --reject-with tcp-reset \
        -m comment --comment "MG: block non-proxy TCP 443"

    # 将自定义链挂载到 OUTPUT（幂等）
    if ! iptables -C OUTPUT -j "$CHAIN_NAME" 2>/dev/null; then
        iptables -A OUTPUT -j "$CHAIN_NAME"
    fi

    # === Rule 4: IPv6 全面阻断 ===
    ip6tables -N "${CHAIN_NAME}_V6" 2>/dev/null || ip6tables -F "${CHAIN_NAME}_V6"
    # 允许回环
    ip6tables -A "${CHAIN_NAME}_V6" -o lo -j ACCEPT \
        -m comment --comment "MG: allow IPv6 loopback"
    # 阻断其他 IPv6 出站
    ip6tables -A "${CHAIN_NAME}_V6" -j DROP \
        -m comment --comment "MG: block all IPv6 outbound"

    if ! ip6tables -C OUTPUT -j "${CHAIN_NAME}_V6" 2>/dev/null; then
        ip6tables -A OUTPUT -j "${CHAIN_NAME}_V6"
    fi

    log "Firewall rules applied successfully."
    log "  - UDP 443/4433: BLOCKED (QUIC)"
    log "  - TCP 443: ONLY '$NODE_PROXY_USER' allowed"
    log "  - IPv6 outbound: BLOCKED"
}

remove_rules() {
    log "Removing fingerprint firewall rules..."

    # 从 OUTPUT 移除引用
    iptables -D OUTPUT -j "$CHAIN_NAME" 2>/dev/null || true
    ip6tables -D OUTPUT -j "${CHAIN_NAME}_V6" 2>/dev/null || true

    # 清空并删除自定义链
    iptables -F "$CHAIN_NAME" 2>/dev/null || true
    iptables -X "$CHAIN_NAME" 2>/dev/null || true
    ip6tables -F "${CHAIN_NAME}_V6" 2>/dev/null || true
    ip6tables -X "${CHAIN_NAME}_V6" 2>/dev/null || true

    log "Firewall rules removed."
}

show_status() {
    log "=== IPv4 MG_FINGERPRINT chain ==="
    iptables -L "$CHAIN_NAME" -n -v 2>/dev/null || echo "(not found)"
    echo
    log "=== IPv6 MG_FINGERPRINT_V6 chain ==="
    ip6tables -L "${CHAIN_NAME}_V6" -n -v 2>/dev/null || echo "(not found)"
}

case "${1:-apply}" in
    apply)  apply_rules ;;
    remove) remove_rules ;;
    status) show_status ;;
    *)
        echo "Usage: $0 [apply|remove|status]"
        exit 1
        ;;
esac