#!/usr/bin/bash

# 获取默认出口网卡（简化版，不含 @ifX）
BASIC_IFACE=$(ip route get 1.1.1.1 | awk '{print $5}')
# 在 `ip link` 里查找完整的 `eth0@ifX`
#FULL_IFACE=$(ip link | awk -v iface="$BASIC_IFACE" '$2 ~ iface {print $2}' | sed 's/://')
ointface=$BASIC_IFACE

# ocserv 配置存在时再读取客户端网段；通用服务器可以没有该配置。
lan=""
if [[ -r /etc/ocserv/ocserv.conf ]]; then
    IPV4_NETWORK=$(awk -F' = ' '/^ipv4-network/ {print $2; exit}' /etc/ocserv/ocserv.conf)
    IPV4_MASK=$(awk -F' = ' '/^ipv4-netmask/ {print $2; exit}' /etc/ocserv/ocserv.conf)
    # 将 `255.255.255.0` 转换为 `/24` 格式
    CIDR=$(awk -F'.' '{for(i=1;i<=NF;i++) {while ($i) {count+=($i%2); $i=int($i/2)}} print count}' <<< "$IPV4_MASK")
    [[ -n "$IPV4_NETWORK" && -n "$CIDR" ]] && lan="$IPV4_NETWORK/$CIDR"
fi

echo "[+] Flushing existing iptables rules..."
iptables -F
iptables -X
iptables -t nat -F
iptables -t nat -X

echo "[+] Disabling IPv4 traffic..."
iptables -P INPUT DROP
#iptables -P OUTPUT DROP
#iptables -P FORWARD DROP

echo "[+] Setting up INPUT chain... "
iptables -A INPUT -m conntrack --ctstate INVALID -j DROP
iptables -A INPUT -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
# 丢弃异常 TCP 新连接，并按来源 IP 限制 SYN 速率
iptables -A INPUT -p tcp ! --syn -m conntrack --ctstate NEW -j DROP
iptables -A INPUT -p tcp --syn -m hashlimit --hashlimit-above 60/second \
    --hashlimit-burst 120 --hashlimit-mode srcip --hashlimit-name tcp_syn -j DROP
iptables -A INPUT -p icmp --icmp-type echo-request -m limit --limit 5/second --limit-burst 10 -j ACCEPT

# 常用网络服务（默认关闭，去掉命令行首的 # 即可启用）
# DNS 服务
echo "[+] Setting up DNS rules..."
#iptables -A INPUT -p tcp --dport 53 -j ACCEPT
#iptables -A INPUT -p udp --dport 53 -j ACCEPT
# FTP 服务；被动端口范围需与 FTP 服务器配置一致
echo "[+] Setting up FTP rules..."
#iptables -A INPUT -p tcp --dport 21 -j ACCEPT
#iptables -A INPUT -p tcp --dport 30000:31000 -j ACCEPT
# NTP 服务
echo "[+] Setting up NTP rules..."
#iptables -A INPUT -p udp --dport 123 -j ACCEPT

# ocserv VPN 服务
echo "[+] Setting up VPN rules..."
iptables -A INPUT -p tcp --dport 6443 -j ACCEPT
iptables -A INPUT -p udp --dport 6443 -j ACCEPT

# Xray / Tunnel 服务
echo "[+] Setting up Xray/Tunnel rules..."
iptables -A INPUT -p tcp --dport 10000:10020 -j ACCEPT

# Web 服务
echo "[+] Setting up WEB rules..."
# HTTP / HTTPS
iptables -A INPUT -p tcp --dport 80 -j ACCEPT
iptables -A INPUT -p tcp --dport 443 -j ACCEPT
# 备用 HTTP / HTTPS
iptables -A INPUT -p tcp --dport 8080 -j ACCEPT
iptables -A INPUT -p tcp --dport 8443 -j ACCEPT
# Grafana（默认关闭）
echo "[+] Setting up Grafana rules..."
#iptables -A INPUT -p tcp --dport 3000 -j ACCEPT
# HTTP Proxy / Squid（默认关闭）
echo "[+] Setting up HTTP Proxy rules..."
#iptables -A INPUT -p tcp --dport 3128 -j ACCEPT
# Prometheus（默认关闭）
echo "[+] Setting up Prometheus rules..."
#iptables -A INPUT -p tcp --dport 9090 -j ACCEPT

# CoDesk 服务
echo "[+] Setting up CoDesk rules..."
# CoDesk 核心和 Web TCP 服务
iptables -A INPUT -p tcp --dport 21114:21119 -j ACCEPT
# CoDesk 集群 TCP 服务
iptables -A INPUT -p tcp -m multiport --dports 21122,21124 -j ACCEPT
# CoDesk UDP P2P 打洞
iptables -A INPUT -p udp --dport 21116 -j ACCEPT

# SSH 远程管理服务
echo "[+] Setting up SSH rules..."
iptables -A INPUT -p tcp --dport 22 -j ACCEPT

# 邮件服务
echo "[+] Setting up EMAIL rules..."
# POP3 / POP3S
#iptables -A INPUT -p tcp --dport 110 -j ACCEPT
#iptables -A INPUT -p tcp --dport 995 -j ACCEPT
# SMTP / SMTPS / Submission
#iptables -A INPUT -p tcp --dport 25 -j ACCEPT
#iptables -A INPUT -p tcp --dport 465 -j ACCEPT
#iptables -A INPUT -p tcp --dport 587 -j ACCEPT
# IMAP / IMAPS
#iptables -A INPUT -p tcp --dport 143 -j ACCEPT
#iptables -A INPUT -p tcp --dport 993 -j ACCEPT

# RDP 远程桌面（默认关闭）
echo "[+] Setting up RDP rules..."
#iptables -A INPUT -p tcp --dport 3389 -j ACCEPT
#iptables -A INPUT -p udp --dport 3389 -j ACCEPT
# WireGuard VPN（默认关闭）
echo "[+] Setting up WireGuard rules..."
#iptables -A INPUT -p udp --dport 51820 -j ACCEPT
# Samba 文件共享（默认关闭，不建议直接暴露到公网）
echo "[+] Setting up Samba rules..."
#iptables -A INPUT -p tcp -m multiport --dports 139,445 -j ACCEPT
# NFS 文件共享（默认关闭，不建议直接暴露到公网）
echo "[+] Setting up NFS rules..."
#iptables -A INPUT -p tcp --dport 2049 -j ACCEPT
#iptables -A INPUT -p udp --dport 2049 -j ACCEPT
# MySQL / MariaDB（默认关闭，建议使用 -s 限制来源）
echo "[+] Setting up MySQL/MariaDB rules..."
#iptables -A INPUT -p tcp --dport 3306 -j ACCEPT
# PostgreSQL（默认关闭，建议使用 -s 限制来源）
echo "[+] Setting up PostgreSQL rules..."
#iptables -A INPUT -p tcp --dport 5432 -j ACCEPT
# Redis（默认关闭，建议使用 -s 限制来源）
echo "[+] Setting up Redis rules..."
#iptables -A INPUT -p tcp --dport 6379 -j ACCEPT
# MongoDB（默认关闭，建议使用 -s 限制来源）
echo "[+] Setting up MongoDB rules..."
#iptables -A INPUT -p tcp --dport 27017 -j ACCEPT

# 允许 VPN 客户端访问本机 DNS；网段从 ocserv 配置自动识别
if [[ -n "$lan" ]]; then
    iptables -A INPUT -p udp -s "$lan" --dport 53 -j ACCEPT
    iptables -A INPUT -p tcp -s "$lan" --dport 53 -j ACCEPT
    # 允许 VPN 客户端访问本机其他 TCP 服务
    iptables -A INPUT -p tcp -s "$lan" -j ACCEPT
fi

echo "[+] ACCEPT traffic from Loopback interface..."
iptables -A INPUT -i lo -j ACCEPT

# LAST RULE - ACCEPT all traffic - Should be commentted once custom rules are created.
#iptables -A INPUT -j ACCEPT

echo "[+] Setting up FORWARD chain..."
iptables -t mangle -I FORWARD 1 -i vpns+ -p tcp --tcp-flags SYN,RST SYN -j TCPMSS --clamp-mss-to-pmtu
iptables -t mangle -I FORWARD 1 -o vpns+ -p tcp --tcp-flags SYN,RST SYN -j TCPMSS --clamp-mss-to-pmtu
iptables -A FORWARD -m conntrack --ctstate INVALID -j DROP
iptables -A FORWARD -m conntrack --ctstate RELATED,ESTABLISHED -j ACCEPT
if [[ -n "$lan" ]]; then
    iptables -A FORWARD -s "$lan" -j ACCEPT
fi
#iptables -A FORWARD -j LOG --log-prefix "IPTABLES-LOG-FORWARD: " --log-level 4
# LAST RULE - ACCEPT all traffic - Should be commentted once custom rules are created.
iptables -A FORWARD -j ACCEPT

echo "[+] Setting up OUTPUT chain..."
#iptables -A OUTPUT -m state --state INVALID -j LOG --log-prefix "IPTABLES-LOG-INVALID:" --log-ip-options --log-tcp-options
iptables -A OUTPUT -m conntrack --ctstate INVALID -j DROP
iptables -A OUTPUT -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
iptables -A OUTPUT -p icmp --icmp-type echo-reply -j ACCEPT
#iptables -A OUTPUT -j LOG --log-prefix "IPTABLES-LOG-OUTPUT: " --log-level 4
# LAST RULE - ACCEPT all traffic - Should be commentted once custom rules are created.
iptables -A OUTPUT -j ACCEPT

# setup the nat
if [[ -n "$lan" ]]; then
    echo "[+] Setting up ocserv NAT rules..."
    iptables -t nat -A POSTROUTING -s "$lan" -o "$ointface" -j MASQUERADE
fi

echo "[+] Save the configuration..."
iptables-save -f /etc/iptables/iptables.rules

echo "[+] Enabling iptables.service at boot..."
if command -v systemctl >/dev/null 2>&1; then
    systemctl enable iptables.service
fi
