#!/bin/bash

# =========================================================================
# WireGuard 统一部署脚本 v3
# 修复：ip_forward 持久化改用 /etc/sysctl.d/99-wg-forward.conf
#       避免被其他 sysctl.d 文件覆盖导致重启后转发失效
# 新增：关闭 rp_filter，避免策略路由/多出口场景丢包
# 新增：PostUp/PostDown 动态获取出口网卡名，不再写死 ens2
# =========================================================================

# --- 基础网络配置 ---
WG_IFACE="wg0"
WG_PORT="1447"
WG_SUBNET="10.100.0"
WG_CONFIG="/etc/wireguard/${WG_IFACE}.conf"

# =========================================================================
# ！！！ 核心密钥配置区 ！！！
# =========================================================================
SERVER_PRI_KEY="MM5+YFr9qH/+1FQaGnmGnZw2i1cQCz2oB2Ab8epclnw="
ROS_CLIENT_PRI_KEY="EPt5s0HSvgeZk48zxPuE4Qf9wj5KZLQNgqUmzFKYGGQ="
PRE_SHARED_KEY="sKXGh2P2JeC5PB0FCofclvk1l76XLdHf7HGrb0VizNo="
TEST_CLIENT_PUB_KEY="MKAmcEQCSuUdw1UY11AjTTnAwM9asAjflm33RB/6Wxs="

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
CYAN='\033[0;36m'
PLAIN='\033[0m'

# =========================================================================
# 工具函数
# =========================================================================
check_root() {
    if [ "$(id -u)" -ne 0 ]; then
        echo -e "${RED}错误：必须以 root 权限运行此脚本。${PLAIN}"
        exit 1
    fi
}

# =========================================================================
# 功能 1：安装 WireGuard
# =========================================================================
install_wg() {
    echo -e "${GREEN}=============================================${PLAIN}"
    echo -e "${GREEN}    开始安装 统一化 WireGuard 服务端...      ${PLAIN}"
    echo -e "${GREEN}=============================================${PLAIN}"

    if [ -z "$SERVER_PRI_KEY" ]; then
        echo -e "${RED}错误：SERVER_PRI_KEY 未填写！${PLAIN}"
        exit 1
    fi

    # -----------------------------------------------------------------
    # [1/7] 安装依赖
    # -----------------------------------------------------------------
    echo -e "${YELLOW}[1/7] 更新系统并安装依赖...${PLAIN}"
    apt-get update -qq
    apt-get install -y wireguard qrencode iptables curl iproute2 openresolv

    # -----------------------------------------------------------------
    # [2/7] 检测网络配置
    # -----------------------------------------------------------------
    echo -e "${YELLOW}[2/7] 检测网络配置...${PLAIN}"
    PUBLIC_IP=$(curl -s --max-time 5 ifconfig.me)
    if [ -z "$PUBLIC_IP" ]; then PUBLIC_IP=$(curl -s --max-time 5 4.ipw.cn); fi
    if [ -z "$PUBLIC_IP" ]; then
        echo -e "${RED}错误：无法获取公网 IP，请检查网络连接。${PLAIN}"
        exit 1
    fi

    PUB_NIC=$(ip -4 route show default | grep -oP '(?<=dev )(\S+)' | head -1)
    if [ -z "$PUB_NIC" ]; then
        echo -e "${RED}错误：无法检测到默认网卡。${PLAIN}"
        exit 1
    fi
    echo -e "公网 IP : ${GREEN}${PUBLIC_IP}${PLAIN}"
    echo -e "主网卡  : ${GREEN}${PUB_NIC}${PLAIN}"

    # -----------------------------------------------------------------
    # [3/7] 生成配置文件
    # -----------------------------------------------------------------
    echo -e "${YELLOW}[3/7] 生成服务器配置 (${WG_CONFIG})...${PLAIN}"
    mkdir -p /etc/wireguard
    chmod 700 /etc/wireguard

    ROS_CLIENT_PUB_KEY=$(echo "${ROS_CLIENT_PRI_KEY}" | wg pubkey)
    SERVER_PUB_KEY=$(echo "${SERVER_PRI_KEY}" | wg pubkey)

    cat > "${WG_CONFIG}" <<EOF
[Interface]
Address = ${WG_SUBNET}.1/24
ListenPort = ${WG_PORT}
PrivateKey = ${SERVER_PRI_KEY}

# 防火墙规则（幂等写法：PostUp 先删后加，避免重启/重载后规则重复堆叠）
# 注意：NIC 变量在 PostUp/PostDown 执行时动态获取当前默认路由出口网卡，
#       不依赖安装时写死的网卡名，适配 VPS 迁移或网卡改名场景。
PostUp   = NIC=\$(ip -4 route show default | grep -oP '(?<=dev )(\S+)' | head -1); \
           iptables -D FORWARD -i ${WG_IFACE} -j ACCEPT 2>/dev/null || true; \
           iptables -D FORWARD -o ${WG_IFACE} -j ACCEPT 2>/dev/null || true; \
           iptables -t nat -D POSTROUTING -o \$NIC -j MASQUERADE 2>/dev/null || true; \
           iptables -I FORWARD -i ${WG_IFACE} -j ACCEPT; \
           iptables -I FORWARD -o ${WG_IFACE} -j ACCEPT; \
           iptables -t nat -A POSTROUTING -o \$NIC -j MASQUERADE
PostDown = NIC=\$(ip -4 route show default | grep -oP '(?<=dev )(\S+)' | head -1); \
           iptables -D FORWARD -i ${WG_IFACE} -j ACCEPT 2>/dev/null || true; \
           iptables -D FORWARD -o ${WG_IFACE} -j ACCEPT 2>/dev/null || true; \
           iptables -t nat -D POSTROUTING -o \$NIC -j MASQUERADE 2>/dev/null || true

# ------------------------------------------
# Peer 1: ROS 路由器 (固定 IP: .2)
# ------------------------------------------
[Peer]
PublicKey = ${ROS_CLIENT_PUB_KEY}
PresharedKey = ${PRE_SHARED_KEY}
AllowedIPs = ${WG_SUBNET}.2/32

# ------------------------------------------
# Peer 2: 测速/连通性节点 test (固定 IP: .3)
# ------------------------------------------
[Peer]
PublicKey = ${TEST_CLIENT_PUB_KEY}
PresharedKey = ${PRE_SHARED_KEY}
AllowedIPs = ${WG_SUBNET}.3/32
EOF
    chmod 600 "${WG_CONFIG}"

    # -----------------------------------------------------------------
    # [4/7] 持久化内核参数 —— 关键修复
    #   ① ip_forward：写到 99-wg-forward.conf，99- 前缀保证加载最靠后，
    #     不被其他 sysctl.d 文件覆盖
    #   ② rp_filter：关闭反向路径过滤，避免策略路由/多出口场景下
    #     内核因回包路径不一致而静默丢包
    # -----------------------------------------------------------------
    echo -e "${YELLOW}[4/7] 持久化内核参数 (ip_forward + rp_filter)...${PLAIN}"
    cat > /etc/sysctl.d/99-wg-forward.conf <<EOF
net.ipv4.ip_forward=1
net.ipv4.conf.all.rp_filter=0
net.ipv4.conf.default.rp_filter=0
EOF
    sysctl -p /etc/sysctl.d/99-wg-forward.conf > /dev/null

    # 立即验证 ip_forward
    CURRENT_FWD=$(cat /proc/sys/net/ipv4/ip_forward)
    if [ "$CURRENT_FWD" != "1" ]; then
        echo -e "${RED}错误：ip_forward 设置失败，当前值=${CURRENT_FWD}，请手动排查！${PLAIN}"
        exit 1
    fi
    echo -e "ip_forward 当前值 : ${GREEN}${CURRENT_FWD}${PLAIN} ✓"
    echo -e "rp_filter  当前值 : ${GREEN}$(cat /proc/sys/net/ipv4/conf/all/rp_filter)${PLAIN} (0=已关闭) ✓"

    # -----------------------------------------------------------------
    # [5/7] 配置 systemd 服务依赖网络就绪后再启动
    #   防止 VPS 重启时 wg-quick 早于网卡就绪导致 PostUp iptables 失败
    # -----------------------------------------------------------------
    echo -e "${YELLOW}[5/7] 配置 wg-quick 服务启动顺序...${PLAIN}"
    mkdir -p /etc/systemd/system/wg-quick@${WG_IFACE}.service.d/
    cat > /etc/systemd/system/wg-quick@${WG_IFACE}.service.d/override.conf <<EOF
[Unit]
After=network-online.target
Wants=network-online.target
EOF
    systemctl daemon-reload

    # -----------------------------------------------------------------
    # [6/7] 启用并启动服务
    # -----------------------------------------------------------------
    echo -e "${YELLOW}[6/7] 启动 WireGuard 服务...${PLAIN}"
    systemctl enable wg-quick@${WG_IFACE}
    systemctl restart wg-quick@${WG_IFACE}

    # -----------------------------------------------------------------
    # [7/7] 验证并输出结果
    # -----------------------------------------------------------------
    if systemctl is-active --quiet wg-quick@${WG_IFACE}; then
        echo -e "${GREEN}[7/7] WireGuard 安装并启动成功！${PLAIN}"
        echo -e "${CYAN}========================================================================${PLAIN}"
        echo -e " 🌐 公网 IP          : ${GREEN}${PUBLIC_IP}${PLAIN}"
        echo -e " 🔑 服务端公钥       : ${GREEN}${SERVER_PUB_KEY}${PLAIN}"
        echo -e " 🖧  WireGuard 端口   : ${GREEN}${WG_PORT}${PLAIN}"
        echo -e " 🚀 ROS 路由器       : ${YELLOW}${WG_SUBNET}.2${PLAIN}（固定）"
        echo -e " 🧪 测速/连通性节点  : ${YELLOW}${WG_SUBNET}.3${PLAIN}（固定）"
        echo -e "${CYAN}========================================================================${PLAIN}"
        echo ""
        echo -e "${YELLOW}验证命令：${PLAIN}"
        echo -e "  wg show wg0                          # 查看连接状态"
        echo -e "  cat /proc/sys/net/ipv4/ip_forward    # 应为 1"
        echo -e "  iptables -t nat -L POSTROUTING -n -v # 应有 MASQUERADE 规则"
    else
        echo -e "${RED}WireGuard 服务启动失败，请检查日志：${PLAIN}"
        echo -e "  journalctl -u wg-quick@${WG_IFACE} --no-pager | tail -30"
        exit 1
    fi
}

# =========================================================================
# 功能 2：添加客户端（手机/电脑，IP 自动顺延从 .4 开始）
# =========================================================================
add_client() {
    if [ ! -f "${WG_CONFIG}" ]; then
        echo -e "${RED}错误：未找到配置文件 ${WG_CONFIG}，请先执行安装！${PLAIN}"
        exit 1
    fi

    read -p "请输入新客户端名称 (如 phone, pc, mac): " CLIENT_NAME
    if [ -z "$CLIENT_NAME" ]; then
        echo -e "${RED}名称不能为空！${PLAIN}"
        exit 1
    fi

    CLIENT_CONFIG_FILE="${HOME}/wg_${CLIENT_NAME}.conf"
    if [ -f "${CLIENT_CONFIG_FILE}" ]; then
        echo -e "${RED}错误：客户端配置 ${CLIENT_CONFIG_FILE} 已存在！${PLAIN}"
        exit 1
    fi

    # 从现有配置里找最大 IP，自动顺延（.2 和 .3 已预留，从 .4 开始）
    LAST_IP=$(grep 'AllowedIPs' "${WG_CONFIG}" | grep -oP '(?<=\.)\d+(?=/32)' | sort -n | tail -1)
    if [ -z "$LAST_IP" ] || [ "$LAST_IP" -lt 3 ]; then
        NEXT_IP=4
    else
        NEXT_IP=$((LAST_IP + 1))
    fi

    if [ "$NEXT_IP" -gt 254 ]; then
        echo -e "${RED}错误：子网 IP 已满（最多 251 个客户端）！${PLAIN}"
        exit 1
    fi
    CLIENT_IP="${WG_SUBNET}.${NEXT_IP}"

    SERVER_PRI_KEY_LOCAL=$(grep '^PrivateKey' "${WG_CONFIG}" | awk '{print $3}')
    SERVER_PUB_KEY=$(echo "${SERVER_PRI_KEY_LOCAL}" | wg pubkey)
    PUBLIC_IP=$(curl -s --max-time 5 ifconfig.me || curl -s --max-time 5 4.ipw.cn)

    CLIENT_PRI_KEY=$(wg genkey)
    CLIENT_PUB_KEY=$(echo "${CLIENT_PRI_KEY}" | wg pubkey)

    echo -e "${YELLOW}为 ${CYAN}${CLIENT_NAME}${YELLOW} 分配 IP: ${CYAN}${CLIENT_IP}${PLAIN}"

    # 追加到服务器配置
    cat >> "${WG_CONFIG}" <<EOF

# Client: ${CLIENT_NAME}
[Peer]
PublicKey = ${CLIENT_PUB_KEY}
PresharedKey = ${PRE_SHARED_KEY}
AllowedIPs = ${CLIENT_IP}/32
EOF

    # 生成客户端配置（AllowedIPs = 0.0.0.0/0 全流量走 VPS）
    cat > "${CLIENT_CONFIG_FILE}" <<EOF
[Interface]
PrivateKey = ${CLIENT_PRI_KEY}
Address = ${CLIENT_IP}/24
DNS = 8.8.8.8, 1.1.1.1

[Peer]
PublicKey = ${SERVER_PUB_KEY}
Endpoint = ${PUBLIC_IP}:${WG_PORT}
PresharedKey = ${PRE_SHARED_KEY}
AllowedIPs = 0.0.0.0/0, ::/0
PersistentKeepalive = 25
EOF
    chmod 600 "${CLIENT_CONFIG_FILE}"

    # 热加载（不重启服务，直接 syncconf 添加 peer）
    wg syncconf ${WG_IFACE} <(wg-quick strip ${WG_IFACE})

    echo -e "${GREEN}=============================================${PLAIN}"
    echo -e "客户端 ${CYAN}${CLIENT_NAME}${PLAIN} 添加成功！IP: ${CYAN}${CLIENT_IP}${PLAIN}"
    echo -e "请使用 WireGuard App 扫描以下二维码："
    echo ""
    qrencode -t ansiutf8 < "${CLIENT_CONFIG_FILE}"
    echo ""
    echo -e "${GREEN}配置文件路径：${CLIENT_CONFIG_FILE}${PLAIN}"
}

# =========================================================================
# 功能 3：重启后自检（快速验证所有关键项）
# =========================================================================
check_status() {
    echo -e "${CYAN}=============================================${PLAIN}"
    echo -e "${CYAN}          WireGuard 状态自检                 ${PLAIN}"
    echo -e "${CYAN}=============================================${PLAIN}"

    # ip_forward
    FWD=$(cat /proc/sys/net/ipv4/ip_forward)
    if [ "$FWD" = "1" ]; then
        echo -e "ip_forward       : ${GREEN}1 ✓${PLAIN}"
    else
        echo -e "ip_forward       : ${RED}0 ✗ （转发未开启！）${PLAIN}"
        echo -e "  修复命令: sysctl -w net.ipv4.ip_forward=1"
    fi

    # WG 服务
    if systemctl is-active --quiet wg-quick@${WG_IFACE}; then
        echo -e "wg-quick@${WG_IFACE}  : ${GREEN}running ✓${PLAIN}"
    else
        echo -e "wg-quick@${WG_IFACE}  : ${RED}stopped ✗${PLAIN}"
        echo -e "  修复命令: systemctl restart wg-quick@${WG_IFACE}"
    fi

    # NAT 规则
    NAT_COUNT=$(iptables -t nat -L POSTROUTING -n -v 2>/dev/null | grep -c MASQUERADE)
    if [ "$NAT_COUNT" -gt 0 ]; then
        echo -e "NAT MASQUERADE   : ${GREEN}存在 (${NAT_COUNT} 条) ✓${PLAIN}"
    else
        echo -e "NAT MASQUERADE   : ${RED}不存在 ✗${PLAIN}"
        echo -e "  修复命令: systemctl restart wg-quick@${WG_IFACE}"
    fi

    # Peers 握手状态
    echo -e "${CYAN}---------------------------------------------${PLAIN}"
    echo -e "Peer 握手状态："
    wg show ${WG_IFACE} 2>/dev/null || echo -e "${RED}wg0 接口不存在${PLAIN}"
}

# =========================================================================
# 交互菜单
# =========================================================================
check_root

echo -e "${CYAN}=============================================${PLAIN}"
echo -e "${CYAN}   WireGuard 管理脚本 v2 (统一密钥版)        ${PLAIN}"
echo -e "${CYAN}=============================================${PLAIN}"
echo -e "  ${GREEN}1.${PLAIN} 安装 WireGuard（全新部署）"
echo -e "  ${GREEN}2.${PLAIN} 添加客户端（手机/PC，IP 自动顺延）"
echo -e "  ${GREEN}3.${PLAIN} 状态自检（验证 ip_forward / NAT / Peer）"
echo -e "  ${GREEN}4.${PLAIN} 退出"
echo -e "${CYAN}=============================================${PLAIN}"
read -p "请输入数字 [1-4]: " MENU_CHOICE

case $MENU_CHOICE in
    1) install_wg ;;
    2) add_client ;;
    3) check_status ;;
    4) exit 0 ;;
    *) echo -e "${RED}输入错误。${PLAIN}"; exit 1 ;;
esac