#!/bin/bash

# 调试输出函数
debug() {
    echo "DEBUG: $*"
}

# 检查是否以 root 用户运行
if [ "$EUID" -ne 0 ]; then
    echo "请以 root 用户运行此脚本"
    exit 1
fi

# 启用 IP 转发
enable_ip_forwarding() {
    local conf_file="/etc/sysctl.conf"
    debug "检查是否可以写入 $conf_file"
    if [ ! -w "$conf_file" ]; then
        echo "无法写入 $conf_file，请检查权限"
        exit 1
    fi

    debug "检查是否已启用 IPv4 转发"
    if grep -q "^net.ipv4.ip_forward=1" "$conf_file"; then
        echo "IPv4 转发已在 $conf_file 中启用"
    else
        echo "net.ipv4.ip_forward=1" >> "$conf_file"
        echo "IPv4 转发已添加到 $conf_file"
    fi

    debug "检查是否已启用 IPv6 转发"
    if grep -q "^net.ipv6.conf.all.forwarding=1" "$conf_file"; then
        echo "IPv6 转发已在 $conf_file 中启用"
    else
        echo "net.ipv6.conf.all.forwarding=1" >> "$conf_file"
        echo "IPv6 转发已添加到 $conf_file"
    fi

    debug "应用 sysctl 配置"
    sysctl -p
}

# 禁用 IP 转发
disable_ip_forwarding() {
    local conf_file="/etc/sysctl.conf"
    debug "检查是否可以写入 $conf_file"
    if [ ! -w "$conf_file" ]; then
        echo "无法写入 $conf_file，请检查权限"
        exit 1
    fi

    debug "禁用 IPv4 和 IPv6 转发"
    sed -i '/^net.ipv4.ip_forward=1/d' "$conf_file"
    sed -i '/^net.ipv6.conf.all.forwarding=1/d' "$conf_file"
    sysctl -w net.ipv4.ip_forward=0
    sysctl -w net.ipv6.conf.all.forwarding=0
    sysctl -p
}

# 自动识别连接到互联网的网卡
get_internet_interface() {
    debug "获取默认网关接口"
    local default_interface
    default_interface=$(ip route | grep '^default' | awk '{print $5}')
    if [ -z "$default_interface" ]; then
        echo "未找到默认网关，无法确定连接互联网的网卡。"
        exit 1
    fi
    echo "$default_interface"
}

# 配置 firewalld 的函数
configure_firewalld() {
    echo "配置 firewalld..."
    local output
    output=$(firewall-cmd --permanent --zone=public --add-masquerade 2>&1)
    if [[ $output == *"ALREADY_ENABLED"* ]]; then
        echo "警告: masquerade 已经启用"
    elif [[ $output == *"Error"* ]]; then
        echo "配置 firewalld 时出错: $output"
        exit 1
    else
        echo "$output"
    fi
    firewall-cmd --reload
}

# 取消 firewalld 配置的函数
unconfigure_firewalld() {
    echo "取消 firewalld 配置..."
    firewall-cmd --permanent --zone=public --remove-masquerade
    firewall-cmd --reload
}

# 配置 iptables 的函数
configure_iptables() {
    local interface=$1
    echo "配置 iptables..."
    if ! iptables -t nat -A POSTROUTING -o "$interface" -j MASQUERADE; then
        echo "配置 iptables 时出错"
        exit 1
    fi

    # 检查是否安装了 iptables-persistent
    if dpkg -l | grep -q iptables-persistent; then
        echo "保存 iptables 规则..."
        if ! netfilter-persistent save; then
            echo "保存 iptables 规则时出错"
            exit 1
        fi
    else
        echo "iptables-persistent 未安装。现在安装..."
        if ! apt-get update || ! apt-get install -y iptables-persistent; then
            echo "安装 iptables-persistent 时出错"
            exit 1
        fi
        if ! netfilter-persistent save; then
            echo "保存 iptables 规则时出错"
            exit 1
        fi
    fi
}

# 取消 iptables 配置的函数
unconfigure_iptables() {
    echo "取消 iptables 配置..."
    iptables -t nat -F
    if dpkg -l | grep -q iptables-persistent; then
        echo "保存 iptables 规则..."
        if ! netfilter-persistent save; then
            echo "保存 iptables 规则时出错"
            exit 1
        fi
    fi
}

# 配置 ufw 的函数
configure_ufw() {
    echo "配置 ufw..."
    if ! ufw status | grep -q "Status: active"; then
        echo "ufw 未启用。现在启用..."
        if ! ufw enable; then
            echo "启用 ufw 时出错"
            exit 1
        fi
    fi

    if ufw status | grep -q "MASQUERADE"; then
        echo "警告: MASQUERADE 已经启用"
    else
        if ! ufw default allow outgoing; then
            echo "设置 ufw 默认策略时出错"
            exit 1
        fi
        if ! ufw default deny incoming; then
            echo "设置 ufw 默认策略时出错"
            exit 1
        fi
        if ! ufw route allow in on any out on any; then
            echo "配置 ufw 转发规则时出错"
            exit 1
        fi
        if ! ufw route allow out on any in on any; then
            echo "配置 ufw 转发规则时出错"
            exit 1
        fi
    fi
}

# 取消 ufw 配置的函数
unconfigure_ufw() {
    echo "取消 ufw 配置..."
    ufw default deny outgoing
    ufw default allow incoming
    ufw disable
}

# 配置端口转发的函数
configure_port_forwarding() {
    local interface=$1
    local protocol=$2
    local ports=$3

    if command -v firewall-cmd &> /dev/null; then
        echo "配置 firewalld 端口转发..."
        for port in ${ports//,/ }; do
            if ! firewall-cmd --permanent --add-forward-port=port=$port:proto=$protocol:toport=$port; then
                echo "配置 firewalld 端口转发时出错"
                exit 1
            fi
        done
        firewall-cmd --reload
    elif command -v iptables &> /dev/null; then
        echo "配置 iptables 端口转发..."
        for port in ${ports//,/ }; do
            if ! iptables -t nat -A PREROUTING -p $protocol --dport $port -j DNAT --to-destination :$port; then
                echo "配置 iptables 端口转发时出错"
                exit 1
            fi
        done
        iptables -t nat -A POSTROUTING -o "$interface" -j MASQUERADE
        if dpkg -l | grep -q iptables-persistent; then
            echo "保存 iptables 规则..."
            if ! netfilter-persistent save; then
                echo "保存 iptables 规则时出错"
                exit 1
            fi
        else
            echo "iptables-persistent 未安装。现在安装..."
            if ! apt-get update || ! apt-get install -y iptables-persistent; then
                echo "安装 iptables-persistent 时出错"
                exit 1
            fi
            if ! netfilter-persistent save; then
                echo "保存 iptables 规则时出错"
                exit 1
            fi
        fi
    elif command -v ufw &> /dev/null; then
        echo "配置 ufw 端口转发..."
        for port in ${ports//,/ }; do
            if ! ufw allow $port/$protocol; then
                echo "配置 ufw 端口转发时出错"
                exit 1
            fi
        done
    else
        echo "未找到支持的防火墙管理工具。请安装 firewalld、iptables 或 ufw。"
        exit 1
    fi
}

# 取消端口转发的函数
unconfigure_all_port_forwarding() {
    if command -v firewall-cmd &> /dev/null; then
        echo "取消 firewalld 端口转发..."
        local ports
        ports=$(firewall-cmd --list-forward-ports)
        for port in ${ports//,/ }; do
            firewall-cmd --permanent --remove-forward-port=$port
        done
        firewall-cmd --reload
    elif command -v iptables &> /dev/null; then
        echo "取消 iptables 端口转发..."
        iptables -t nat -F PREROUTING
        iptables -t nat -F POSTROUTING
        if dpkg -l | grep -q iptables-persistent; then
            echo "保存 iptables 规则..."
            if ! netfilter-persistent save; then
                echo "保存 iptables 规则时出错"
                exit 1
            fi
        fi
    elif command -v ufw &> /dev/null; then
        echo "取消 ufw 端口转发..."
        ufw reset
    else
        echo "未找到支持的防火墙管理工具。请安装 firewalld、iptables 或 ufw。"
        exit 1
    fi
}


# 获取有效的整数输入
get_valid_integer_input() {
    local prompt=$1
    local input

    while true; do
        read -p "$prompt" input
        if [[ "$input" =~ ^[0-9]+$ ]]; then
            echo "$input"
            return
        else
            echo "无效的输入，请输入一个整数。"
        fi
    done
}

# 获取有效的端口输入
get_valid_ports_input() {
    local prompt=$1
    local input

    while true; do
        read -p "$prompt" input
        if [[ "$input" =~ ^[0-9]+(,[0-9]+)*$ ]]; then
            echo "$input"
            return
        else
            echo "无效的输入，请输入一个或多个端口号，用逗号分隔。"
        fi
    done
}

# 显示当前配置状态
show_current_configuration() {
    echo "当前配置状态:"
    echo "----------------------------"
    echo "IP 转发状态:"
    sysctl net.ipv4.ip_forward
    sysctl net.ipv6.conf.all.forwarding
    echo "----------------------------"
    echo "防火墙状态:"
    if command -v firewall-cmd &> /dev/null; then
        firewall-cmd --list-all
    elif command -v iptables &> /dev/null; then
        iptables -t nat -L -n -v
    elif command -v ufw &> /dev/null; then
        ufw status
    else
        echo "未找到支持的防火墙管理工具。"
    fi
    echo "----------------------------"
}

# 主函数
main() {
    show_current_configuration

    echo "请选择操作类型:"
    echo "1. 配置全局转发或端口转发"
    echo "2. 取消配置并还原"
    echo "3. 退出"
    choice=$(get_valid_integer_input "请输入选项 (1, 2 或 3): ")

    if [ "$choice" -eq 1 ]; then
        debug "启用 IP 转发"
        enable_ip_forwarding

        # 获取连接互联网的网卡
        debug "获取连接互联网的网卡"
        local internet_interface
        internet_interface=$(get_internet_interface)
        echo "检测到的互联网网卡: $internet_interface"

        echo "请选择操作类型:"
        echo "1. 全局转发"
        echo "2. 指定端口转发"
        echo "3. 取消配置"
        sub_choice=$(get_valid_integer_input "请输入选项 (1, 2 或 3): ")

        if [ "$sub_choice" -eq 1 ]; then
            debug "配置全局转发"
            # 全局转发
            if command -v firewall-cmd &> /dev/null; then
                configure_firewalld
            elif command -v iptables &> /dev/null; then
                configure_iptables "$internet_interface"
            elif command -v ufw &> /dev/null; then
                configure_ufw
            else
                echo "未找到支持的防火墙管理工具。请安装 firewalld、iptables 或 ufw。"
                exit 1
            fi
        elif [ "$sub_choice" -eq 2 ]; then
            debug "配置指定端口转发"
            # 指定端口转发
            echo "请选择协议类型:"
            echo "1. TCP"
            echo "2. UDP"
            protocol_choice=$(get_valid_integer_input "请输入选项 (1 或 2): ")

            if [ "$protocol_choice" -eq 1 ]; then
                protocol="tcp"
            elif [ "$protocol_choice" -eq 2 ]; then
                protocol="udp"
            else
                echo "无效的选项"
                exit 1
            fi

            echo "请选择端口类型:"
            echo "1. 单个端口"
            echo "2. 连续端口"
            echo "3. 不连续端口"
            port_choice=$(get_valid_integer_input "请输入选项 (1, 2 或 3): ")

            if [ "$port_choice" -eq 1 ]; then
                ports=$(get_valid_integer_input "请输入端口号: ")
            elif [ "$port_choice" -eq 2 ]; then
                start_port=$(get_valid_integer_input "请输入起始端口号: ")
                end_port=$(get_valid_integer_input "请输入结束端口号: ")
                ports=$(seq -s , $start_port $end_port)
            elif [ "$port_choice" -eq 3 ]; then
                ports=$(get_valid_ports_input "请输入端口号（用逗号分隔）: ")
            else
                echo "无效的选项"
                exit 1
            fi

            configure_port_forwarding "$internet_interface" "$protocol" "$ports"
        elif [ "$sub_choice" -eq 3 ]; then
            debug "取消配置"
            disable_ip_forwarding

            if command -v firewall-cmd &> /dev/null; then
                unconfigure_firewalld
            elif command -v iptables &> /dev/null; then
                unconfigure_iptables
            elif command -v ufw &> /dev/null; then
                unconfigure_ufw
            else
                echo "未找到支持的防火墙管理工具。请安装 firewalld、iptables 或 ufw。"
                exit 1
            fi

            unconfigure_all_port_forwarding

            echo "已取消配置并还原。"
        else
            echo "无效的选项"
            exit 1
        fi

        echo "IP 转发和防火墙配置完成。"

    elif [ "$choice" -eq 2 ]; then
        debug "取消所有配置"
        disable_ip_forwarding

        if command -v firewall-cmd &> /dev/null; then
            unconfigure_firewalld
        elif command -v iptables &> /dev/null; then
            unconfigure_iptables
        elif command -v ufw &> /dev/null; then
            unconfigure_ufw
        else
            echo "未找到支持的防火墙管理工具。请安装 firewalld、iptables 或 ufw。"
            exit 1
        fi

        unconfigure_all_port_forwarding

        echo "已取消配置并还原。"

    elif [ "$choice" -eq 3 ]; then
        echo "退出脚本。"
        exit 0
    else
        echo "无效的选项"
        exit 1
    fi
}

# 执行主函数
main