#!/bin/bash

# 日志记录
log() {
    local level=$1
    local message=$2
    local timestamp=$(date +"%Y-%m-%d %H:%M:%S")
    local log_file="/var/log/firewall_manager.log"
    
    # 创建日志文件如果不存在
    [ ! -f "$log_file" ] && touch "$log_file"
    
    # 设置日志级别颜色
    case $level in
        INFO) color="\033[32m" ;;
        WARN) color="\033[33m" ;;
        ERROR) color="\033[31m" ;;
        *) color="\033[0m" ;;
    esac
    
    # 输出到控制台和日志文件
    echo -e "${color}[$timestamp] [$level] $message\033[0m"
    echo "[$timestamp] [$level] $message" >> "$log_file"
}

# 检测系统类型
detect_system() {
    if [ -f /etc/centos-release ]; then
        echo "centos"
    elif [ -f /etc/lsb-release ]; then
        echo "ubuntu"
    elif [ -f /etc/openwrt_release ]; then
        echo "openwrt"
    elif [ -f /etc/debian_version ]; then
        echo "debian"
    else
        log "ERROR" "无法检测系统类型"
        exit 1
    fi
}

# 检测防火墙状态
check_firewall_status() {
    log "INFO" "========== 防火墙状态 =========="
    
    # 显示防火墙运行状态
    case $SYSTEM in
        centos)
            if systemctl status firewalld &> /dev/null; then
                log "INFO" "防火墙状态：已启用"
            else
                log "WARN" "防火墙状态：未启用"
            fi
            ;;
        ubuntu|debian)
            if ufw status | grep -q "Status: active"; then
                log "INFO" "防火墙状态：已启用"
            else
                log "WARN" "防火墙状态：未启用"
            fi
            ;;
        openwrt)
            if /etc/init.d/firewall status &> /dev/null; then
                log "INFO" "防火墙状态：已启用"
            else
                log "WARN" "防火墙状态：未启用"
            fi
            ;;
        *)
            log "ERROR" "不支持的系统类型"
            exit 1
            ;;
    esac
    
    # 显示已开放端口
    log "INFO" ""
    log "INFO" "========== 已开放端口 =========="
    list_ports
    
    # 显示转发规则
    log "INFO" ""
    log "INFO" "========== 转发规则 =========="
    list_forward_rules
    
    # 显示IP转发状态
    log "INFO" ""
    log "INFO" "========== IP转发状态 =========="
    show_forward_status
    
    log "INFO" ""
    log "INFO" "==============================="
    return 0
}

# 检测当前正在使用的端口
detect_active_ports() {
    case $SYSTEM in
        centos|ubuntu|debian)
            netstat -tuln | awk '/^tcp|^udp/ {split($4, a, ":"); print a[length(a)]}' | sort -u
            ;;
        openwrt)
            netstat -tuln | awk '/^tcp|^udp/ {split($4, a, ":"); print a[length(a)]}' | sort -u
            ;;
    esac
}

# 开启防火墙
enable_firewall() {
    log "INFO" "正在检测当前使用的端口..."
    ports=$(detect_active_ports)
    
    if [ -n "$ports" ]; then
        log "INFO" "检测到以下正在使用的端口："
        log "INFO" "$ports"
        log "INFO" "正在将这些端口添加到防火墙规则中..."
        
        for port in $ports; do
            case $SYSTEM in
                centos)
                    firewall-cmd --permanent --add-port=$port/tcp
                    firewall-cmd --permanent --add-port=$port/udp
                    ;;
                ubuntu|debian)
                    ufw allow $port/tcp
                    ufw allow $port/udp
                    ;;
                openwrt)
                    iptables -A INPUT -p tcp --dport $port -j ACCEPT
                    iptables -A INPUT -p udp --dport $port -j ACCEPT
                    /etc/init.d/firewall save
                    ;;
            esac
            log "INFO" "端口 $port 已添加"
        done
    else
        log "INFO" "未检测到正在使用的端口"
    fi
    
    case $SYSTEM in
        centos)
            systemctl start firewalld
            systemctl enable firewalld
            firewall-cmd --reload
            ;;
        ubuntu|debian)
            ufw enable
            ;;
        openwrt)
            /etc/init.d/firewall start
            /etc/init.d/firewall enable
            ;;
    esac
    log "INFO" "防火墙已成功启用"
}

# 查看已开放端口
list_ports() {
    case $SYSTEM in
        centos)
            firewall-cmd --list-ports
            ;;
        ubuntu|debian)
            ufw status | grep "ALLOW"
            ;;
        openwrt)
            iptables -L -n | grep "ACCEPT"
            ;;
    esac
}

# 添加端口
add_port() {
    local port=$1
    local protocol=${2:-tcp}
    
    case $SYSTEM in
        centos)
            firewall-cmd --permanent --add-port=$port/$protocol
            firewall-cmd --reload
            ;;
        ubuntu|debian)
            ufw allow $port/$protocol
            ;;
        openwrt)
            iptables -A INPUT -p $protocol --dport $port -j ACCEPT
            /etc/init.d/firewall save
            ;;
    esac
    log "INFO" "端口 $port/$protocol 已添加"
}

# 删除端口
remove_port() {
    local port=$1
    local protocol=${2:-tcp}
    
    log "INFO" "当前开放的端口："
    list_ports
    
    case $SYSTEM in
        centos)
            firewall-cmd --permanent --remove-port=$port/$protocol
            firewall-cmd --reload
            ;;
        ubuntu|debian)
            ufw delete allow $port/$protocol
            ;;
        openwrt)
            iptables -D INPUT -p $protocol --dport $port -j ACCEPT
            /etc/init.d/firewall save
            ;;
    esac
    log "INFO" "端口 $port/$protocol 已移除"
}

# 启用IP转发
enable_ip_forward() {
    case $SYSTEM in
        centos)
            sysctl -w net.ipv4.ip_forward=1
            echo "net.ipv4.ip_forward=1" > /etc/sysctl.d/99-ipforward.conf
            ;;
        ubuntu|debian)
            sysctl -w net.ipv4.ip_forward=1
            sed -i 's/#*net.ipv4.ip_forward=.*/net.ipv4.ip_forward=1/' /etc/sysctl.conf
            ;;
        openwrt)
            uci set network.lan.ipv4_forward=1
            uci commit network
            /etc/init.d/network restart
            ;;
    esac
    log "INFO" "IP转发已启用"
}

# 禁用IP转发
disable_ip_forward() {
    case $SYSTEM in
        centos)
            sysctl -w net.ipv4.ip_forward=0
            echo "net.ipv4.ip_forward=0" > /etc/sysctl.d/99-ipforward.conf
            ;;
        ubuntu|debian)
            sysctl -w net.ipv4.ip_forward=0
            sed -i 's/#*net.ipv4.ip_forward=.*/net.ipv4.ip_forward=0/' /etc/sysctl.conf
            ;;
        openwrt)
            uci set network.lan.ipv4_forward=0
            uci commit network
            /etc/init.d/network restart
            ;;
    esac
    log "INFO" "IP转发已禁用"
}

# 查看转发清单
list_forward_rules() {
    log "INFO" "当前转发规则："
    case $SYSTEM in
        centos)
            firewall-cmd --list-forward-ports
            ;;
        ubuntu|debian)
            iptables -t nat -L -n -v | grep "DNAT"
            ;;
        openwrt)
            uci show firewall | grep redirect
            ;;
    esac
}

# 查看转发状态
show_forward_status() {
    case $SYSTEM in
        centos|ubuntu|debian)
            sysctl net.ipv4.ip_forward
            ;;
        openwrt)
            uci get network.lan.ipv4_forward
            ;;
    esac
}

# 基本转发（单个端口）
add_single_forward() {
    local src_port=$1
    local dest_ip=$2
    local dest_port=$3
    local protocol=${4:-tcp}
    
    case $SYSTEM in
        centos)
            firewall-cmd --permanent --add-forward-port=port=$src_port:proto=$protocol:toport=$dest_port:toaddr=$dest_ip
            firewall-cmd --reload
            ;;
        ubuntu|debian)
            iptables -t nat -A PREROUTING -p $protocol --dport $src_port -j DNAT --to-destination $dest_ip:$dest_port
            iptables -t nat -A POSTROUTING -p $protocol -d $dest_ip --dport $dest_port -j MASQUERADE
            ;;
        openwrt)
            uci add firewall redirect
            uci set firewall.@redirect[-1].src='wan'
            uci set firewall.@redirect[-1].src_dport=$src_port
            uci set firewall.@redirect[-1].dest_ip=$dest_ip
            uci set firewall.@redirect[-1].dest_port=$dest_port
            uci set firewall.@redirect[-1].proto=$protocol
            uci commit firewall
            /etc/init.d/firewall restart
            ;;
    esac
    log "INFO" "基本转发已添加：$src_port -> $dest_ip:$dest_port ($protocol)"
}

# 全局转发（所有流量）
add_global_forward() {
    local dest_ip=$1
    
    case $SYSTEM in
        centos)
            firewall-cmd --permanent --add-masquerade
            firewall-cmd --permanent --add-rich-rule="rule family=ipv4 source address=0.0.0.0/0 destination address=$dest_ip forward-port port=1-65535 protocol=tcp to-port=1-65535"
            firewall-cmd --permanent --add-rich-rule="rule family=ipv4 source address=0.0.0.0/0 destination address=$dest_ip forward-port port=1-65535 protocol=udp to-port=1-65535"
            firewall-cmd --reload
            ;;
        ubuntu|debian)
            iptables -t nat -A PREROUTING -j DNAT --to-destination $dest_ip
            iptables -t nat -A POSTROUTING -j MASQUERADE
            ;;
        openwrt)
            uci set firewall.@redirect[0].src='wan'
            uci set firewall.@redirect[0].dest_ip=$dest_ip
            uci commit firewall
            /etc/init.d/firewall restart
            ;;
    esac
    log "INFO" "全局转发已添加：所有流量 -> $dest_ip"
}

# 端口范围转发
add_range_forward() {
    local start_port=$1
    local end_port=$2
    local dest_ip=$3
    local dest_start_port=$4
    local protocol=${5:-tcp}
    
    case $SYSTEM in
        centos)
            for ((port=start_port; port<=end_port; port++)); do
                dest_port=$((dest_start_port + port - start_port))
                firewall-cmd --permanent --add-forward-port=port=$port:proto=$protocol:toport=$dest_port:toaddr=$dest_ip
            done
            firewall-cmd --reload
            ;;
        ubuntu|debian)
            for ((port=start_port; port<=end_port; port++)); do
                dest_port=$((dest_start_port + port - start_port))
                iptables -t nat -A PREROUTING -p $protocol --dport $port -j DNAT --to-destination $dest_ip:$dest_port
                iptables -t nat -A POSTROUTING -p $protocol -d $dest_ip --dport $dest_port -j MASQUERADE
            done
            ;;
        openwrt)
            for ((port=start_port; port<=end_port; port++)); do
                dest_port=$((dest_start_port + port - start_port))
                uci add firewall redirect
                uci set firewall.@redirect[-1].src='wan'
                uci set firewall.@redirect[-1].src_dport=$port
                uci set firewall.@redirect[-1].dest_ip=$dest_ip
                uci set firewall.@redirect[-1].dest_port=$dest_port
                uci set firewall.@redirect[-1].proto=$protocol
            done
            uci commit firewall
            /etc/init.d/firewall restart
            ;;
    esac
    log "INFO" "端口范围转发已添加：$start_port-$end_port -> $dest_ip:$dest_start_port-$(($dest_start_port + $end_port - $start_port)) ($protocol)"
}

# 查看已有转发策略
list_forward_rules() {
    case $SYSTEM in
        centos)
            firewall-cmd --list-forward-ports
            ;;
        ubuntu|debian)
            iptables -t nat -L -n | grep "DNAT"
            ;;
        openwrt)
            uci show firewall | grep redirect
            ;;
    esac
}

# 指定源IP转发
add_source_forward() {
    local src_ip=$1
    local src_port=$2
    local dest_ip=$3
    local dest_port=$4
    local protocol=${5:-tcp}
    
    case $SYSTEM in
        centos)
            firewall-cmd --permanent --add-rich-rule="rule family=ipv4 source address=$src_ip forward-port port=$src_port protocol=$protocol to-port=$dest_port to-addr=$dest_ip"
            firewall-cmd --reload
            ;;
        ubuntu|debian)
            iptables -t nat -A PREROUTING -s $src_ip -p $protocol --dport $src_port -j DNAT --to-destination $dest_ip:$dest_port
            iptables -t nat -A POSTROUTING -s $src_ip -p $protocol -d $dest_ip --dport $dest_port -j MASQUERADE
            ;;
        openwrt)
            uci add firewall redirect
            uci set firewall.@redirect[-1].src='wan'
            uci set firewall.@redirect[-1].src_ip=$src_ip
            uci set firewall.@redirect[-1].src_dport=$src_port
            uci set firewall.@redirect[-1].dest_ip=$dest_ip
            uci set firewall.@redirect[-1].dest_port=$dest_port
            uci set firewall.@redirect[-1].proto=$protocol
            uci commit firewall
            /etc/init.d/firewall restart
            ;;
    esac
    log "INFO" "指定源转发已添加：$src_ip:$src_port -> $dest_ip:$dest_port ($protocol)"
}

# 删除转发规则
remove_forward() {
    log "INFO" "========== 当前转发规则 =========="
    
    # 获取转发规则并编号
    case $SYSTEM in
        centos)
            rules=($(firewall-cmd --list-forward-ports))
            ;;
        ubuntu|debian)
            rules=($(iptables -t nat -L -n -v | grep "DNAT" | awk '{print $11,$12,$13}'))
            ;;
        openwrt)
            rules=($(uci show firewall | grep redirect | awk -F= '{print $2}'))
            ;;
    esac
    
    # 显示带编号的规则列表
    if [ ${#rules[@]} -eq 0 ]; then
        log "INFO" "当前没有转发规则"
        log "INFO" ""
        log "INFO" "提示：您可以通过菜单选项8来添加新的转发规则"
        log "INFO" "==============================="
        return 0
    fi
    
    for i in "${!rules[@]}"; do
        log "INFO" "$((i+1)). ${rules[$i]}"
    done
    echo ""
    
    # 让用户选择要删除的规则
    read -p "请输入要删除的规则编号（1-${#rules[@]}）：" rule_num
    if ! [[ "$rule_num" =~ ^[0-9]+$ ]] || [ "$rule_num" -lt 1 ] || [ "$rule_num" -gt "${#rules[@]}" ]; then
        log "ERROR" "无效的规则编号"
        return 1
    fi
    
    # 获取选中的规则
    selected_rule=${rules[$((rule_num-1))]}
    
    log "INFO" "您将要删除以下转发规则："
    log "INFO" "$selected_rule"
    
    read -p "确认删除？(y/n)：" confirm
    if [[ "$confirm" != "y" && "$confirm" != "Y" ]]; then
        log "INFO" "取消删除操作"
        return 0
    fi
    
    # 执行删除操作
    case $SYSTEM in
        centos)
            firewall-cmd --permanent --remove-forward-port=$selected_rule
            firewall-cmd --reload
            ;;
        ubuntu|debian)
            iptables -t nat -D PREROUTING -p $(echo $selected_rule | awk '{print $1}' | awk -F: '{print $1}') --dport $(echo $selected_rule | awk '{print $1}' | awk -F: '{print $2}') -j DNAT
            iptables -t nat -D POSTROUTING -p $(echo $selected_rule | awk '{print $1}' | awk -F: '{print $1}') -d $(echo $selected_rule | awk '{print $3}') -j MASQUERADE
            ;;
        openwrt)
            uci delete $(uci show firewall | grep "$selected_rule" | cut -d'=' -f1 | sed 's/\.name$//')
            uci commit firewall
            /etc/init.d/firewall restart
            ;;
    esac
    
    log "INFO" "转发规则已删除：$selected_rule"
    log "INFO" ""
    log "INFO" "========== 更新后的转发规则 =========="
    list_forward_rules
}

# 显示菜单
show_menu() {
    echo ""
    echo "=============================="
    echo "  防火墙管理菜单"
    echo "=============================="
    echo "1. 检查防火墙状态"
    echo "2. 启用防火墙"
    echo "3. 查看开放端口"
    echo "4. 添加新端口"
    echo "5. 删除端口"
    echo "6. 启用IP转发"
    echo "7. 禁用IP转发"
    echo "8. 添加端口转发"
    echo "9. 删除端口转发"
    echo "10. 查看转发规则"
    echo "0. 退出"
    echo "=============================="
    echo ""
    echo "请输入选项编号并按回车："
}

# 环境检测
detect_environment() {
    log "INFO" "========== 环境检测开始 =========="
    
    # 系统信息
    SYSTEM=$(detect_system)
    log "INFO" "系统类型：$SYSTEM"
    
    # 系统版本
    case $SYSTEM in
        centos)
            os_version=$(cat /etc/centos-release)
            ;;
        ubuntu)
            os_version=$(lsb_release -d | awk -F"\t" '{print $2}')
            ;;
        debian)
            os_version=$(cat /etc/debian_version)
            ;;
        openwrt)
            os_version=$(cat /etc/openwrt_release | grep DISTRIB_DESCRIPTION | awk -F"'" '{print $2}')
            ;;
        *)
            os_version="未知"
            ;;
    esac
    log "INFO" "系统版本：$os_version"
    
    # 内核版本
    kernel_version=$(uname -r)
    log "INFO" "内核版本：$kernel_version"
    
    # 网络信息
    log "INFO" ""
    log "INFO" "========== 网络信息 =========="
    ip_addr=$(ip -4 addr show | grep inet | awk '{print $2}')
    log "INFO" "IP地址："
    log "INFO" "$ip_addr"
    
    # 防火墙状态
    log "INFO" ""
    log "INFO" "========== 防火墙状态 =========="
    check_firewall_status
    
    log "INFO" "==============================="
}

# 主程序
main() {
    # 执行环境检测
    detect_environment
    
    # 检测系统类型
    SYSTEM=$(detect_system)
    
    while true; do
        show_menu
        read -p "请输入选项编号：" choice
        
        case $choice in
            1)
                if check_firewall_status; then
                    log "INFO" "防火墙已启用"
                else
                    log "WARN" "防火墙未启用"
                fi
                ;;
            2)
                enable_firewall
                ;;
            3)
                list_ports
                ;;
            4)
                read -p "请输入要添加的端口号：" port
                read -p "请输入协议类型（tcp/udp，默认为tcp）：" protocol
                add_port $port ${protocol:-tcp}
                ;;
            5)
                read -p "请输入要删除的端口号：" port
                read -p "请输入协议类型（tcp/udp，默认为tcp）：" protocol
                remove_port $port ${protocol:-tcp}
                ;;
            6)
                enable_ip_forward
                ;;
            7)
                disable_ip_forward
                ;;
            8)
                read -p "请输入源端口号：" src_port
                read -p "请输入目标IP地址：" dest_ip
                read -p "请输入目标端口号：" dest_port
                read -p "请输入协议类型（tcp/udp，默认为tcp）：" protocol
                add_single_forward $src_port $dest_ip $dest_port ${protocol:-tcp}
                ;;
            9)
                read -p "请输入要删除的源端口号：" src_port
                read -p "请输入协议类型（tcp/udp，默认为tcp）：" protocol
                remove_forward $src_port ${protocol:-tcp}
                ;;
            10)
                list_forward_rules
                ;;
            0)
                log "INFO" "退出防火墙管理程序"
                exit 0
                ;;
            *)
                log "ERROR" "无效的选项，请重新选择"
                ;;
        esac
        
        read -p "按回车键继续..."
    done
}

# 启动主程序
main
