#!/bin/bash
# wg-manager.sh - WireGuard 跨平台管理脚本 (使用 192.168.108.0/24 网段)

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

# 变量定义 - 修改为 192.168.108.0/24 网段
CONFIG_DIR="/etc/wireguard"
SERVER_CONFIG="$CONFIG_DIR/wg0.conf"
SERVER_PUBLIC_KEY="$CONFIG_DIR/server_public.key"
SERVER_PRIVATE_KEY="$CONFIG_DIR/server_private.key"
CLIENTS_DIR="$CONFIG_DIR/clients"
NETWORK="192.168.108.0/24"
SERVER_IP="192.168.108.1"
PORT="51820"

# 日志函数
log() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')]${NC} $1"
}

warn() {
    echo -e "${YELLOW}[警告]${NC} $1"
}

error() {
    echo -e "${RED}[错误]${NC} $1"
    exit 1
}

# 检测操作系统
detect_os() {
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        OS=$ID
        OS_VERSION=$VERSION_ID
    else
        OS=$(uname -s)
    fi
    log "检测到操作系统: $OS $OS_VERSION"
}

# 安装 WireGuard
install_wireguard() {
    log "开始安装 WireGuard..."
    
    case $OS in
        ubuntu|debian)
            sudo apt update || warn "apt update 失败，继续安装..."
            sudo apt install -y wireguard-tools resolvconf qrencode || {
                warn "仓库安装失败，尝试编译安装..."
                compile_install
            }
            ;;
        centos|rhel|fedora|rocky|almalinux|opencloudos)
            if command -v dnf >/dev/null 2>&1; then
                sudo dnf install -y wireguard-tools qrencode || {
                    sudo dnf install -y epel-release
                    sudo dnf install -y wireguard-tools qrencode
                }
            else
                sudo yum install -y wireguard-tools qrencode || {
                    sudo yum install -y epel-release
                    sudo yum install -y wireguard-tools qrencode
                }
            fi
            ;;
        *)
            warn "不支持的操作系统 $OS，尝试编译安装..."
            compile_install
            ;;
    esac
    
    # 验证安装
    if command -v wg >/dev/null 2>&1; then
        log "WireGuard 安装成功: $(wg --version)"
    else
        error "WireGuard 安装失败"
    fi
}

# 编译安装
compile_install() {
    log "开始编译安装 WireGuard..."
    
    # 安装编译依赖
    case $OS in
        ubuntu|debian)
            sudo apt install -y git build-essential libmnl-dev pkg-config
            ;;
        *)
            if command -v dnf >/dev/null 2>&1; then
                sudo dnf groupinstall -y "Development Tools"
                sudo dnf install -y git libmnl-devel pkgconfig gcc-c++
            elif command -v yum >/dev/null 2>&1; then
                sudo yum groupinstall -y "Development Tools"
                sudo yum install -y git libmnl-devel pkgconfig gcc-c++
            fi
            ;;
    esac
    
    # 编译安装 wireguard-tools
    cd /tmp
    rm -rf wireguard-tools
    git clone https://git.zx2c4.com/wireguard-tools
    cd wireguard-tools/src
    make
    sudo make install
    
    # 安装 qrencode（可选）
    if command -v apt >/dev/null 2>&1; then
        sudo apt install -y qrencode
    elif command -v dnf >/dev/null 2>&1; then
        sudo dnf install -y qrencode
    elif command -v yum >/dev/null 2>&1; then
        sudo yum install -y qrencode
    fi
}

# 启用 IP 转发
enable_ip_forward() {
    log "启用 IP 转发..."
    
    # 临时启用
    echo 1 | sudo tee /proc/sys/net/ipv4/ip_forward >/dev/null
    
    # 永久启用
    if grep -q "net.ipv4.ip_forward" /etc/sysctl.conf; then
        sudo sed -i 's/^#*net.ipv4.ip_forward=.*/net.ipv4.ip_forward=1/' /etc/sysctl.conf
    else
        echo 'net.ipv4.ip_forward=1' | sudo tee -a /etc/sysctl.conf
    fi
    
    sudo sysctl -p >/dev/null 2>&1
    log "IP 转发已启用"
}

# 安全的文件操作函数
create_secure_file() {
    local file_path="$1"
    local content="$2"
    
    # 创建目录
    sudo mkdir -p "$(dirname "$file_path")"
    
    # 创建临时文件
    local temp_file=$(mktemp)
    echo "$content" > "$temp_file"
    
    # 复制并设置权限
    sudo cp "$temp_file" "$file_path"
    sudo chmod 600 "$file_path"
    sudo chown root:root "$file_path"
    
    # 清理临时文件
    rm -f "$temp_file"
}

# 生成随机客户端IP (192.168.108.2 - 192.168.108.254)
generate_client_ip() {
    local base_ip="192.168.108"
    local random_suffix=$(( RANDOM % 253 + 2 ))  # 2-254
    echo "${base_ip}.${random_suffix}"
}

# 初始化服务端
init_server() {
    log "初始化 WireGuard 服务端..."
    log "使用网段: $NETWORK"
    log "服务器IP: $SERVER_IP"
    
    sudo mkdir -p $CONFIG_DIR
    sudo mkdir -p $CLIENTS_DIR
    
    # 生成服务器密钥
    if [ ! -f "$SERVER_PRIVATE_KEY" ]; then
        log "生成服务器密钥对..."
        sudo wg genkey | sudo tee $SERVER_PRIVATE_KEY | sudo wg pubkey | sudo tee $SERVER_PUBLIC_KEY >/dev/null
        sudo chmod 600 $SERVER_PRIVATE_KEY $SERVER_PUBLIC_KEY
    fi
    
    # 获取默认网卡
    DEFAULT_IFACE=$(ip route | grep default | awk '{print $5}' | head -1)
    [ -z "$DEFAULT_IFACE" ] && DEFAULT_IFACE="eth0"
    
    # 创建服务端配置
    create_secure_file "$SERVER_CONFIG" "
[Interface]
PrivateKey = $(sudo cat $SERVER_PRIVATE_KEY)
Address = $SERVER_IP/24
ListenPort = $PORT
SaveConfig = false

# 流量转发规则
PostUp = iptables -A FORWARD -i wg0 -j ACCEPT; iptables -t nat -A POSTROUTING -o $DEFAULT_IFACE -j MASQUERADE
PostDown = iptables -D FORWARD -i wg0 -j ACCEPT; iptables -t nat -D POSTROUTING -o $DEFAULT_IFACE -j MASQUERADE
"

    log "服务端配置已创建: $SERVER_CONFIG"
    log "服务器公钥: $(sudo cat $SERVER_PUBLIC_KEY)"
}

# 启动服务端
start_server() {
    log "启动 WireGuard 服务端..."
    
    sudo systemctl enable wg-quick@wg0
    if sudo systemctl is-active wg-quick@wg0 >/dev/null 2>&1; then
        sudo systemctl restart wg-quick@wg0
    else
        sudo systemctl start wg-quick@wg0
    fi
    
    sleep 2
    if sudo systemctl is-active wg-quick@wg0 >/dev/null 2>&1; then
        log "WireGuard 服务端启动成功"
    else
        error "WireGuard 服务端启动失败"
    fi
}

# 添加客户端
add_client() {
    local client_name=$1
    local client_ip=$2
    
    [ -z "$client_ip" ] && client_ip=$(generate_client_ip)
    
    log "添加客户端: $client_name (IP: $client_ip)"
    
    # 生成客户端密钥
    local client_private_key=$(wg genkey)
    local client_public_key=$(echo $client_private_key | wg pubkey)
    
    # 获取服务器公网IP（用于客户端配置）
    local server_public_ip=$(curl -s -4 ifconfig.me || curl -s -4 ipinfo.io/ip || echo "YOUR_SERVER_IP")
    
    # 创建客户端配置
    local client_config="$CLIENTS_DIR/$client_name.conf"
    local client_config_content="[Interface]
PrivateKey = $client_private_key
Address = $client_ip/24
DNS = 8.8.8.8

[Peer]
PublicKey = $(sudo cat $SERVER_PUBLIC_KEY)
Endpoint = $server_public_ip:$PORT
AllowedIPs = 192.168.108.0/24
PersistentKeepalive = 25"
    
    create_secure_file "$client_config" "$client_config_content"
    
    # 添加到服务端配置
    sudo wg set wg0 peer "$client_public_key" allowed-ips "$client_ip/32"
    
    # 保存到服务端配置文件中
    echo "
# Client: $client_name
[Peer]
PublicKey = $client_public_key
AllowedIPs = $client_ip/32" | sudo tee -a $SERVER_CONFIG > /dev/null

    # 生成二维码（如果安装了 qrencode）
    if command -v qrencode >/dev/null 2>&1; then
        sudo qrencode -t ansiutf8 < "$client_config"
        sudo qrencode -o "$CLIENTS_DIR/$client_name.png" < "$client_config"
        log "客户端二维码已生成: $CLIENTS_DIR/$client_name.png"
    fi
    
    log "客户端添加成功!"
    echo "=========================================="
    echo "客户端名称: $client_name"
    echo "客户端IP: $client_ip"
    echo "客户端公钥: $client_public_key"
    echo "配置文件: $client_config"
    echo "=========================================="
    
    # 显示客户端配置
    echo -e "\n客户端配置内容:"
    sudo cat "$client_config"
}

# 批量添加内网节点
add_internal_nodes() {
    log "开始批量添加内网节点..."
    log "使用网段: $NETWORK"
    
    local nodes=("node1" "node2" "node3" "node4")
    local start_ip=2  # 从 192.168.108.2 开始
    
    for i in "${!nodes[@]}"; do
        local client_name="${nodes[$i]}"
        local client_ip="192.168.108.$((start_ip + i))"
        log "添加节点: $client_name -> $client_ip"
        add_client "$client_name" "$client_ip"
        echo
    done
    
    log "内网节点批量添加完成"
    log "节点IP分配:"
    for i in "${!nodes[@]}"; do
        echo "  ${nodes[$i]} -> 192.168.108.$((start_ip + i))"
    done
}

# 显示状态
show_status() {
    log "WireGuard 状态:"
    sudo wg show
    
    echo -e "\n连接的客户端:"
    sudo wg show | grep -A5 "peer" | grep -E "peer:|allowed ips:" | sed 's/peer:/客户端:/;s/allowed ips:/IP地址:/'
    
    echo -e "\nVPN 网络信息:"
    echo "网段: $NETWORK"
    echo "服务器: $SERVER_IP"
    echo "端口: $PORT"
    
    echo -e "\n配置文件列表:"
    sudo ls -la $CLIENTS_DIR/ 2>/dev/null || warn "客户端配置目录不存在"
}

# 显示使用说明
show_usage() {
    echo -e "${BLUE}WireGuard 管理脚本 (网段: $NETWORK)${NC}"
    echo "用法: $0 [选项]"
    echo
    echo "选项:"
    echo "  install        安装 WireGuard 并启用 IP 转发"
    echo "  server         初始化并启动服务端"
    echo "  add-client <名称> [IP]  添加客户端"
    echo "  add-nodes      批量添加内网节点 (node1-node4)"
    echo "  status         显示状态"
    echo "  restart        重启服务"
    echo "  full-setup     完整安装配置（安装+服务端+内网节点）"
    echo
    echo "示例:"
    echo "  $0 install           # 安装 WireGuard"
    echo "  $0 full-setup        # 完整安装配置"
    echo "  $0 add-client mypc   # 添加客户端"
    echo "  $0 status            # 查看状态"
}

# 完整安装配置
full_setup() {
    log "开始完整安装配置 WireGuard..."
    log "VPN 网段: $NETWORK"
    detect_os
    install_wireguard
    enable_ip_forward
    init_server
    start_server
    add_internal_nodes
    show_status
    log "完整安装配置完成!"
    
    echo -e "\n${GREEN}下一步操作:${NC}"
    echo "1. 在云服务器安全组开放 $PORT UDP 端口"
    echo "2. 将客户端配置文件分发到对应节点"
    echo "3. 在客户端安装 WireGuard 并导入配置"
    echo "4. 使用 VPN IP 建立 Docker Swarm 集群"
    echo -e "\n${YELLOW}重要信息:${NC}"
    echo "服务器公钥: $(sudo cat $SERVER_PUBLIC_KEY)"
    echo "VPN 网段: $NETWORK"
    echo "服务器内网IP: $SERVER_IP"
}

# 主函数
main() {
    case $1 in
        install)
            detect_os
            install_wireguard
            enable_ip_forward
            ;;
        server)
            init_server
            start_server
            ;;
        add-client)
            if [ -z "$2" ]; then
                error "请提供客户端名称"
            fi
            add_client "$2" "$3"
            ;;
        add-nodes)
            add_internal_nodes
            ;;
        status)
            show_status
            ;;
        restart)
            sudo systemctl restart wg-quick@wg0
            log "服务已重启"
            ;;
        full-setup)
            full_setup
            ;;
        *)
            show_usage
            ;;
    esac
}

# 脚本入口
if [ $# -eq 0 ]; then
    show_usage
else
    main "$@"
fi