#!/bin/bash

install_wireguard() {
    echo "开始安装WireGuard..."
    # Check Linux distribution
    if [[ -f /etc/redhat-release ]]; then
        # Install WireGuard on CentOS / RHEL
        yum install epel-release -y
        yum install wireguard-tools -y
    elif [[ -f /etc/debian_version ]]; then
        # Install WireGuard on Debian / Ubuntu
        apt update
        apt install wireguard -y
    elif [[ -f /etc/arch-release ]]; then
        # Install WireGuard on Arch Linux
        pacman -S wireguard-tools --noconfirm
    elif [[ "$(uname)" == "Darwin" ]]; then
        # Install WireGuard on macOS (using Homebrew)
        brew install wireguard-tools
    else
        echo "不受支持的 Linux 发行版或操作系统"
        exit 1
    fi

    # Start the WireGuard service
    systemctl enable wg-quick@wg0
    systemctl start wg-quick@wg0

    # Check if WireGuard was installed successfully
    if ! command -v wg &>/dev/null; then
        echo "WireGuard安装失败"
        exit 1
    fi

    echo "WireGuard安装成功"
}

generate_keys() {
    cd /etc/wireguard
    # 生成WireGuard私钥和公钥
    echo "正在生成WireGuard私钥和公钥..."
    wg genkey | tee privatekey | wg pubkey > publickey && cat privatekey && cat publickey
}

create_server_config() {
    # 服务器端配置
    read -p "请输入服务器的公网IP地址： " server_public_ip
    read -p "请输入服务器的内网IP地址： " server_internal_ip
    read -p "请输入WireGuard服务器端口号： " server_port
    read -p "请输入客户端的公钥： " client_public_key
    read -p "请输入客户端的内网IP地址： " client_internal_ip
    eth=$(ls /sys/class/net| grep ^e | head -n1)

    # 创建服务器端配置文件
    echo "[Interface]" > /etc/wireguard/wg0.conf
    echo "PrivateKey = $(cat /etc/wireguard/privatekey)" >> /etc/wireguard/wg0.conf
    echo "Address = $server_internal_ip/24" >> /etc/wireguard/wg0.conf
    echo "PostUp   = iptables -A FORWARD -i wg0 -j ACCEPT; iptables -A FORWARD -o wg0 -j ACCEPT; iptables -t nat -A POSTROUTING -o $eth -j MASQUERADE" >> /etc/wireguard/wg0.conf
    echo "PostDown = iptables -D FORWARD -i wg0 -j ACCEPT; iptables -D FORWARD -o wg0 -j ACCEPT; iptables -t nat -D POSTROUTING -o $eth -j MASQUERADE" >> /etc/wireguard/wg0.conf
    echo "ListenPort = $server_port" >> /etc/wireguard/wg0.conf

    echo "" >> /etc/wireguard/wg0.conf
    echo "[Peer]" >> /etc/wireguard/wg0.conf
    echo "PublicKey = $client_public_key" >> /etc/wireguard/wg0.conf
    echo "AllowedIPs = $client_internal_ip/32" >> /etc/wireguard/wg0.conf

    echo "WireGuard服务器端配置文件已成功创建在 /etc/wireguard/wg0.conf。"
}

create_client_config() {
    # 客户端配置
    read -p "请输入服务器的公网IP地址： " server_public_ip
    read -p "请输入WireGuard服务器端口号： " server_port
    read -p "请输入服务器的公钥： " server_public_key
    read -p "请输入客户端的内网IP地址： " client_internal_ip
    eth=$(ls /sys/class/net| grep ^e | head -n1)

    # 创建客户端配置文件
    echo "[Interface]" > /etc/wireguard/wg0.conf
    echo "PrivateKey = $(cat /etc/wireguard/privatekey)" >> /etc/wireguard/wg0.conf
    echo "Address = $client_internal_ip/24" >> /etc/wireguard/wg0.conf

    echo "" >> /etc/wireguard/wg0.conf
    echo "[Peer]" >> /etc/wireguard/wg0.conf
    echo "PublicKey = $server_public_key" >> /etc/wireguard/wg0.conf
    echo "Endpoint = $server_public_ip:$server_port" >> /etc/wireguard/wg0.conf
    echo "AllowedIPs = 0.0.0.0/0, ::/0" >> /etc/wireguard/wg0.conf

    echo "WireGuard客户端配置文件已成功创建在 /etc/wireguard/wg0.conf。"
}

start_wireguard() {
    echo "正在启动WireGuard..."
    wg-quick up wg0
    echo "WireGuard已启动"
}

stop_wireguard() {
    echo "正在停止WireGuard..."
    wg-quick down wg0
    echo "WireGuard已停止"
}

restart_wireguard() {
    echo "正在重启WireGuard..."
    systemctl restart wg-quick@wg0
    echo "WireGuard已重启"
}

main() {
    # 显示菜单并选择角色
    echo -e "---------------------------"
    echo -e "欢迎使用WireGuard配置脚本"
    echo -e "\t---authored by yanglc---"
    echo -e "---------------------------"
    echo "0. 退出脚本"
    echo "1. 安装Wireguard"
    echo "2. 配置服务器端（中转）"
    echo "3. 配置客户端（落地）"
    echo "4. 启动WireGuard"
    echo "5. 停止WireGuard"
    echo "6. 重启WireGuard（systemctl）"
    echo "7. 重启Wireguard（wg-quick）"
    read -p "请输入数字: " role

    case $role in
        0)
            exit 0
            ;;
        1)
            install_wireguard
            generate_keys
            main
            ;;
        2)
            create_server_config
            main
            ;;
        3)
            create_client_config
            main
            ;;
        4)
            start_wireguard
            main
            ;;
        5)
            stop_wireguard
            main
            ;;
        6)
            restart_wireguard
            main
            ;;
        7)
            stop_wireguard
            start_wireguard
            main
            ;;
        *)
            echo "输入无效的选项。"
            main
            ;;
    esac
}

main