#!/bin/bash

# 确保脚本以 root 权限运行
if [[ $EUID -ne 0 ]]; then
   echo "错误：本脚本必须以 root 权限运行" 
   exit 1
fi

# 检查是否安装了 nftables
if ! command -v nft &> /dev/null; then
    echo "未检查到 nftables，正在尝试安装..."
    if command -v apt-get &> /dev/null; then
        apt-get update -y && apt-get install -y nftables
    elif command -v yum &> /dev/null; then
        yum install -y nftables
    else
        echo "错误：无法自动安装 nftables，请手动安装后重试。"
        exit 1
    fi
    systemctl enable nftables --now
fi

# 初始化 nftables 的 nat 表和链
function init_nftables() {
    nft list table ip nat &>/dev/null
    if [ $? -ne 0 ]; then
        echo ""
        read -p "首次初始化，请输入外部网络接口名称 (oifname, 例如 eth0): " eth_name
        if [[ -z "$eth_name" ]]; then
            eth_name="eth0"
            echo "未输入，默认使用 eth0"
        fi

        nft add table ip nat
        # prerouting 链用于 DNAT (修改目标地址)
        nft add chain ip nat prerouting { type nat hook prerouting priority 0 \; policy accept \; }
        # postrouting 链用于 SNAT/masquerade (修改源地址)
        nft add chain ip nat postrouting { type nat hook postrouting priority 100 \; policy accept \; }
        
        # 添加通用 masquerade 规则
        nft add rule ip nat postrouting oifname "$eth_name" masquerade
    fi
}

function save_rules() {
    echo "正在保存 nftables 规则..."
    if [ -f /etc/redhat-release ]; then
        nft list ruleset > /etc/sysconfig/nftables.conf || echo "保存失败"
    else
        nft list ruleset > /etc/nftables.conf || echo "保存失败"
    fi
    echo "规则保存完毕。"
}

function add_rule() {
    init_nftables

    echo ""
    read -p "请输入本地监听端口 (例如 8080): " local_port
    read -p "请输入目标 IP 地址 (例如 10.0.0.2): " dest_ip
    read -p "请输入目标端口 (不输入则默认和本地监听端口一致): " dest_port
    read -p "请输入转发协议 (tcp/udp/both) [默认: both]: " protocol
    protocol=${protocol:-both}
    
    if [[ -z "$dest_port" ]]; then
        dest_port=$local_port
    fi

    if [[ -z "$local_port" || -z "$dest_ip" ]]; then
        echo "本地端口和目标IP不能为空，操作取消。"
        return
    fi

    echo "正在添加规则..."

    if [[ "$protocol" == "tcp" || "$protocol" == "both" ]]; then
        if [[ "$local_port" == "$dest_port" ]]; then
            nft add rule ip nat prerouting tcp dport $local_port dnat to $dest_ip
        else
            nft add rule ip nat prerouting tcp dport $local_port dnat to $dest_ip:$dest_port
        fi
    fi

    if [[ "$protocol" == "udp" || "$protocol" == "both" ]]; then
        if [[ "$local_port" == "$dest_port" ]]; then
            nft add rule ip nat prerouting udp dport $local_port dnat to $dest_ip
        else
            nft add rule ip nat prerouting udp dport $local_port dnat to $dest_ip:$dest_port
        fi
    fi

    echo -e "\n成功添加端口转发规则: 本机:$local_port -> $dest_ip:$dest_port (协议: $protocol)"
    save_rules
}

function list_rules() {
    echo ""
    echo "========== 当前的 NAT 转发规则 =========="
    nft -a list table ip nat 2>/dev/null || echo "当前没有任何 NAT 规则。"
    echo "========================================="
}

function del_rule() {
    list_rules
    echo ""
    echo "提示: 删除规则需要提供上方输出中的链名称 (chain) 和句柄编号 (handle)。"
    read -p "请输入所在链的名称 (例如 prerouting 或 postrouting): " chain_name
    read -p "请输入要删除的 rule handle 编号: " handle_num
    
    if [[ -n "$chain_name" && -n "$handle_num" ]]; then
        nft delete rule ip nat $chain_name handle $handle_num
        if [ $? -eq 0 ]; then
            echo "规则 (handle $handle_num) 已删除。"
            save_rules
        else
            echo "错误: 删除失败，请检查链名和句柄编号是否正确。"
        fi
    else
        echo "输入无效，操作取消。"
    fi
}

function main() {
    while true; do
        echo ""
        echo "==================================="
        echo "    nftables 端口转发管理脚本       "
        echo "==================================="
        echo "1. 添加端口转发规则"
        echo "2. 查看当前转发规则"
        echo "3. 删除特定转发规则"
        echo "4. 退出脚本"
        echo "==================================="
        read -p "请选择一个操作 [1-4]: " option

        case $option in
            1) add_rule ;;
            2) list_rules ;;
            3) del_rule ;;
            4) 
                echo "退出脚本..."
                exit 0 
                ;;
            *) 
                echo "无效的选项，请重新选择。" 
                ;;
        esac
    done
}

main