#!/bin/bash

# Copyright (C) 2006 OpenWrt.org
# Copyright 2022-2026 sirpdboy <herboy2008@gmail.com>

crrun=$1
crid=$2
NAME=timecontrol
DEBUG=1  # 开启调试

config_t_get() {
    local index=${3:-0}
    local ret=$(uci -q get "${NAME}.@${1}[${index}].${2}")
    echo "${ret:-$4}"
}

LOG_FILE="/var/log/timecontrol.log"
IDLIST="/var/$NAME.idlist"

bin_nft=$(which nft 2>/dev/null)
bin_iptables=$(which iptables 2>/dev/null)
bin_ip6tables=$(which ip6tables 2>/dev/null)
bin_conntrack=$(which conntrack 2>/dev/null)

nftables_ver=0
iptables_ver=0

# 获取配置
chain=$(config_t_get timecontrol chain 0 "forward")
list_type=$(config_t_get timecontrol list_type 0 "blacklist")

if [ "$chain" = "input" ]; then
    StrongCHAIN=1
else
    StrongCHAIN=0
fi

dbg() {
    if [ "$DEBUG" -eq 1 ]; then
        local d="$(date '+%Y-%m-%d %H:%M:%S')"
        echo "[$d] FW-DEBUG: $@" >> "$LOG_FILE"
        echo "FW-DEBUG: $@"
    fi
}

info() {
    local d="$(date '+%Y-%m-%d %H:%M:%S')"
    echo "[$d] FW-INFO: $@" >> "$LOG_FILE"
    echo "FW-INFO: $@"
}

# 地址解析函数 - 修复格式问题
parse_target() {
    local target="$1"
    
    # 去除空格
    target=$(echo "${target}" | xargs)
    
    # dbg "解析目标地址: $target"
    
    # IPv4单个地址
    if echo "$target" | grep -qE '^([0-9]{1,3}\.){3}[0-9]{1,3}$'; then
        local octets=(${target//./ })
        local valid=1
        for octet in "${octets[@]}"; do
            if [ "$octet" -gt 255 ] || [ "$octet" -lt 0 ]; then
                valid=0
                break
            fi
        done
        [ "$valid" -eq 1 ] && {
            echo "ipv4:single:$target"
            return 0
        }
    
    # IPv4范围
    elif echo "$target" | grep -qE '^([0-9]{1,3}\.){3}[0-9]{1,3}-([0-9]{1,3}\.){3}[0-9]{1,3}$'; then
        local start_ip=${target%-*}
        local end_ip=${target#*-}
        echo "ipv4:range:$start_ip-$end_ip"
        return 0
    
    # CIDR
    elif echo "$target" | grep -qE '^([0-9]{1,3}\.){3}[0-9]{1,3}/[0-9]{1,2}$'; then
        local ip=${target%/*}
        local mask=${target#*/}
        [ "$mask" -le 32 ] && [ "$mask" -ge 0 ] && {
            echo "ipv4:cidr:$target"
            return 0
        }
    
    # MAC地址
    elif echo "$target" | grep -qE '^([0-9A-Fa-f]{2}:){5}[0-9A-Fa-f]{2}$'; then
        echo "mac:single:$(echo "$target" | tr '[:upper:]' '[:lower:]')"
        return 0
    
    # IPv6地址
    elif echo "$target" | grep -qE '^([0-9a-fA-F]{0,4}:){2,7}[0-9a-fA-F]{0,4}$'; then
        echo "ipv6:single:$target"
        return 0
    
    # IPv6 CIDR
    elif echo "$target" | grep -qE '^([0-9a-fA-F]{0,4}:){2,7}[0-9a-fA-F]{0,4}/[0-9]{1,3}$'; then
        local ipv6=${target%/*}
        local mask=${target#*/}
        [ "$mask" -le 128 ] && [ "$mask" -ge 0 ] && {
            echo "ipv6:cidr:$target"
            return 0
        }
    fi
    
    dbg "无法解析地址: $target"
    return 1
}

# 清理现有连接
flush_connections() {
    local target="$1"
    
    [ -x "$bin_conntrack" ] || {
        dbg "conntrack不可用"
        return
    }
    
    local parsed_result=$(parse_target "$target")
    [ $? -eq 0 ] || {
        dbg "无法解析地址用于清理连接: $target"
        return
    }
    
    IFS=':' read -r type subtype value <<< "$parsed_result"
    
    dbg "清理连接: type=$type, value=$value"
    
    case "$type" in
        "ipv4")
            $bin_conntrack -D -s "$value" 2>/dev/null && dbg "清理源连接: $value"
            $bin_conntrack -D -d "$value" 2>/dev/null && dbg "清理目标连接: $value"
            ;;
        "mac")
            # MAC地址需要先转换为IP
            if [ -f "/proc/net/arp" ]; then
                local ip_addr=$(grep -i "$value" /proc/net/arp 2>/dev/null | awk '{print $1}' | head -1)
                if [ -n "$ip_addr" ]; then
                    $bin_conntrack -D -s "$ip_addr" 2>/dev/null && dbg "清理MAC源连接: $value -> $ip_addr"
                    $bin_conntrack -D -d "$ip_addr" 2>/dev/null && dbg "清理MAC目标连接: $value -> $ip_addr"
                fi
            fi
            ;;
    esac
}

# 检查防火墙工具
check_firewall_tool() {
    if [ -x "$bin_nft" ]; then
        nftables_ver=1
        dbg "检测到nftables: $bin_nft"
    elif [ -x "$bin_iptables" ] && [ -x "$bin_ip6tables" ]; then
        iptables_ver=1
        dbg "检测到iptables: $bin_iptables, $bin_ip6tables"
    else
        info "错误: 未找到可用的防火墙工具"
        return 1
    fi
    return 0
}

# 初始化防火墙
init_firewall() {
    check_firewall_tool || return 1
    
    info "初始化防火墙规则 (模式: $list_type, 强度: $chain)"
    
    if [ -n "$nftables_ver" ]; then
        # 使用nftables
        dbg "初始化nftables"
        
        # 删除可能存在的旧表
        nft delete table inet timecontrol 2>/dev/null
        sleep 1
        
        # 创建新表
        nft add table inet timecontrol
        nft add chain inet timecontrol forward "{ type filter hook forward priority -100; policy accept; }"
        
        # 创建黑名单集合
        nft add set inet timecontrol blacklist "{ type ipv4_addr; flags interval; }"
        nft add set inet timecontrol blacklist6 "{ type ipv6_addr; flags interval; }"
        nft add set inet timecontrol blacklist_mac "{ type ether_addr; }"
        
        # 添加规则（黑名单模式：匹配到就DROP）
        nft add rule inet timecontrol forward ip saddr @blacklist drop
        nft add rule inet timecontrol forward ip6 saddr @blacklist6 drop
        nft add rule inet timecontrol forward ether saddr @blacklist_mac drop
        
        # 强控制模式
        if [ "$StrongCHAIN" -eq 1 ]; then
            nft add chain inet timecontrol input "{ type filter hook input priority -100; policy accept; }"
            nft add rule inet timecontrol input ip saddr @blacklist drop
            nft add rule inet timecontrol input ip6 saddr @blacklist6 drop
            nft add rule inet timecontrol input ether saddr @blacklist_mac drop
            dbg "已启用强控制模式 (INPUT链)"
        fi
        
        info "nftables初始化完成"
        
    elif [ -n "$iptables_ver" ]; then
        # 使用iptables
        dbg "初始化iptables"
        
        # 创建ipset（如果不存在）
        ipset create timecontrol_blacklist hash:net 2>/dev/null || {
            ipset flush timecontrol_blacklist
            dbg "已存在的ipset timecontrol_blacklist已清空"
        }
        
        ipset create timecontrol_blacklist6 hash:net family inet6 2>/dev/null || {
            ipset flush timecontrol_blacklist6
            dbg "已存在的ipset timecontrol_blacklist6已清空"
        }
        
        # 删除可能存在的旧规则
        iptables -D FORWARD -m set --match-set timecontrol_blacklist src -j DROP 2>/dev/null
        ip6tables -D FORWARD -m set --match-set timecontrol_blacklist6 src -j DROP 2>/dev/null
        
        # 添加新规则（黑名单模式）
        iptables -I FORWARD -m set --match-set timecontrol_blacklist src -j DROP
        ip6tables -I FORWARD -m set --match-set timecontrol_blacklist6 src -j DROP
        
        dbg "已添加FORWARD规则"
        
        # 强控制模式
        if [ "$StrongCHAIN" -eq 1 ]; then
            iptables -D INPUT -m set --match-set timecontrol_blacklist src -j DROP 2>/dev/null
            ip6tables -D INPUT -m set --match-set timecontrol_blacklist6 src -j DROP 2>/dev/null
            
            iptables -I INPUT -m set --match-set timecontrol_blacklist src -j DROP
            ip6tables -I INPUT -m set --match-set timecontrol_blacklist6 src -j DROP
            
            dbg "已启用强控制模式 (INPUT链)"
        fi
        
        info "iptables初始化完成"
    fi
    
    return 0
}

# 停止防火墙规则
stop_firewall() {
    info "停止防火墙规则"
    
    if [ -n "$nftables_ver" ]; then
        nft delete table inet timecontrol 2>/dev/null && info "nftables规则已删除"
    fi
    
    if [ -n "$iptables_ver" ]; then
        # 删除iptables规则
        iptables -D FORWARD -m set --match-set timecontrol_blacklist src -j DROP 2>/dev/null
        iptables -D INPUT -m set --match-set timecontrol_blacklist src -j DROP 2>/dev/null
        ip6tables -D FORWARD -m set --match-set timecontrol_blacklist6 src -j DROP 2>/dev/null
        ip6tables -D INPUT -m set --match-set timecontrol_blacklist6 src -j DROP 2>/dev/null
        
        # 删除ipset
        ipset destroy timecontrol_blacklist 2>/dev/null
        ipset destroy timecontrol_blacklist6 2>/dev/null
        
        info "iptables规则已删除"
    fi
    
    # 清理ID列表
    rm -f "$IDLIST"
}

# 添加设备到防火墙
add_device() {
    local id="$1"
    local target=$(config_t_get device mac "$id")
    [ -z "$target" ] && {
        dbg "添加设备失败: ID $id 的目标地址为空"
        return
    }
    
    local comment=$(config_t_get device comment "$id" "设备$id")
    info "添加设备到防火墙: $comment ($target)"
    
    local parsed_result=$(parse_target "$target")
    if [ $? -ne 0 ]; then
        info "添加失败: 无法解析地址 $target"
        return
    fi
    
    IFS=':' read -r type subtype value <<< "$parsed_result"
    dbg "解析结果: type=$type, subtype=$subtype, value=$value"
    
    if [ -n "$nftables_ver" ]; then
        # nftables处理
        case "$type" in
            "ipv4")
                nft add element inet timecontrol blacklist "{ $value }" 2>&1 | while read line; do dbg "nft: $line"; done
                dbg "已添加到nftables黑名单(IPv4): $value"
                ;;
            "ipv6")
                nft add element inet timecontrol blacklist6 "{ $value }" 2>&1 | while read line; do dbg "nft: $line"; done
                dbg "已添加到nftables黑名单(IPv6): $value"
                ;;
            "mac")
                nft add element inet timecontrol blacklist_mac "{ $value }" 2>&1 | while read line; do dbg "nft: $line"; done
                dbg "已添加到nftables黑名单(MAC): $value"
                ;;
        esac
        
    elif [ -n "$iptables_ver" ]; then
        # iptables处理
        case "$type" in
            "ipv4")
                ipset add timecontrol_blacklist "$value" 2>&1 | while read line; do dbg "ipset: $line"; done
                dbg "已添加到ipset黑名单(IPv4): $value"
                ;;
            "ipv6")
                ipset add timecontrol_blacklist6 "$value" 2>&1 | while read line; do dbg "ipset: $line"; done
                dbg "已添加到ipset黑名单(IPv6): $value"
                ;;
            "mac")
                # iptables不支持MAC地址直接过滤，记录日志
                info "警告: iptables不支持MAC地址过滤，设备 $target 可能无法被阻止"
                ;;
        esac
    fi
    
    # 强控制模式清理连接
    if [ "$StrongCHAIN" -eq 1 ]; then
        dbg "强控制模式，清理现有连接"
        flush_connections "$target"
    fi
    
    # 验证规则
    verify_firewall_rule "$target"
}

# 验证防火墙规则
verify_firewall_rule() {
    local target="$1"
    
    dbg "验证防火墙规则: $target"
    
    if [ -n "$nftables_ver" ]; then
        nft list table inet timecontrol 2>/dev/null | grep -q "$target" && {
            dbg "验证成功: $target 在nftables规则中"
            return 0
        }
    elif [ -n "$iptables_ver" ]; then
        ipset test timecontrol_blacklist "$target" 2>/dev/null && {
            dbg "验证成功: $target 在ipset中"
            return 0
        }
    fi
    
    dbg "验证失败: $target 不在防火墙规则中"
    return 1
}

# 从防火墙移除设备
del_device() {
    local id="$1"
    local target=$(config_t_get device mac "$id")
    [ -z "$target" ] && {
        dbg "移除设备失败: ID $id 的目标地址为空"
        return
    }
    
    local comment=$(config_t_get device comment "$id" "设备$id")
    info "从防火墙移除设备: $comment ($target)"
    
    local parsed_result=$(parse_target "$target")
    [ $? -eq 0 ] || {
        info "移除失败: 无法解析地址 $target"
        return
    }
    
    IFS=':' read -r type subtype value <<< "$parsed_result"
    
    if [ -n "$nftables_ver" ]; then
        case "$type" in
            "ipv4")
                nft delete element inet timecontrol blacklist "{ $value }" 2>/dev/null
                dbg "已从nftables移除(IPv4): $value"
                ;;
            "ipv6")
                nft delete element inet timecontrol blacklist6 "{ $value }" 2>/dev/null
                dbg "已从nftables移除(IPv6): $value"
                ;;
            "mac")
                nft delete element inet timecontrol blacklist_mac "{ $value }" 2>/dev/null
                dbg "已从nftables移除(MAC): $value"
                ;;
        esac
        
    elif [ -n "$iptables_ver" ]; then
        case "$type" in
            "ipv4")
                ipset del timecontrol_blacklist "$value" 2>/dev/null
                dbg "已从ipset移除(IPv4): $value"
                ;;
            "ipv6")
                ipset del timecontrol_blacklist6 "$value" 2>/dev/null
                dbg "已从ipset移除(IPv6): $value"
                ;;
        esac
    fi
}

# 显示防火墙状态
show_firewall_status() {
    echo ""
    echo "防火墙状态:"
    echo "控制模式: $list_type"
    echo "控制强度: $chain $( [ "$StrongCHAIN" -eq 1 ] && echo "(强控制)" )"
    echo ""
    
    if [ -n "$nftables_ver" ]; then
        echo "nftables规则:"
        nft list table inet timecontrol 2>/dev/null || echo "  未找到timecontrol表"
    elif [ -n "$iptables_ver" ]; then
        echo "iptables规则:"
        echo "FORWARD链:"
        iptables -L FORWARD -n | grep -i timecontrol || echo "  未找到timecontrol规则"
        ip6tables -L FORWARD -n | grep -i timecontrol || echo "  未找到IPv6 timecontrol规则"
        
        if [ "$StrongCHAIN" -eq 1 ]; then
            echo ""
            echo "INPUT链:"
            iptables -L INPUT -n | grep -i timecontrol || echo "  未找到timecontrol规则"
            ip6tables -L INPUT -n | grep -i timecontrol || echo "  未找到IPv6 timecontrol规则"
        fi
        
        echo ""
        echo "ipset内容:"
        ipset list timecontrol_blacklist 2>/dev/null | head -20 || echo "  timecontrol_blacklist未找到"
        echo ""
        ipset list timecontrol_blacklist6 2>/dev/null | head -20 || echo "  timecontrol_blacklist6未找到"
    fi
}

# 诊断函数
diagnose() {
    echo ""
    echo "=== 时间控制系统诊断 ==="
    echo ""
    
    # 检查服务
    echo "1. 服务状态:"
    if ps | grep -q "timecontrolctrl"; then
        echo "   ✓ timecontrolctrl 正在运行"
    else
        echo "   ✗ timecontrolctrl 未运行"
    fi
    
    # 检查配置文件
    echo ""
    echo "2. 配置文件:"
    if [ -f "/etc/config/timecontrol" ]; then
        echo "   ✓ 配置文件存在"
        uci show timecontrol 2>/dev/null | grep -c "device" | while read count; do
            echo "   配置了 $count 个设备"
        done
    else
        echo "   ✗ 配置文件不存在"
    fi
    
    # 检查防火墙工具
    echo ""
    echo "3. 防火墙工具:"
    if [ -x "$bin_nft" ]; then
        echo "   ✓ nftables: $bin_nft"
        echo "   版本: $($bin_nft --version 2>/dev/null | head -1)"
    elif [ -x "$bin_iptables" ]; then
        echo "   ✓ iptables: $bin_iptables"
        echo "   版本: $($bin_iptables --version 2>/dev/null | head -1)"
    else
        echo "   ✗ 未找到防火墙工具"
    fi
    
    # 显示当前规则
    show_firewall_status
    
    # 检查ID列表
    echo ""
    echo "4. 当前控制列表:"
    if [ -f "$IDLIST" ] && [ -s "$IDLIST" ]; then
        echo "   当前禁止的设备:"
        cat "$IDLIST" | sed 's/!//g' | while read id; do
            local target=$(config_t_get device mac "$id")
            local comment=$(config_t_get device comment "$id" "设备$id")
            echo "   ID$id: $comment ($target)"
        done
    else
        echo "   当前没有设备被禁止"
    fi
    
    echo ""
    echo "=== 诊断完成 ==="
}

# 主命令处理
case "$crrun" in
    "start")
        info "启动时间控制"
        stop_firewall
        init_firewall
        if [ $? -eq 0 ]; then
            info "时间控制启动成功"
            show_firewall_status
        else
            info "时间控制启动失败"
        fi
        ;;
        
    "stop")
        info "停止时间控制"
        stop_firewall
        info "时间控制已停止"
        ;;
        
    "add")
        [ -z "$crid" ] && {
            echo "错误: 需要指定设备ID"
            exit 1
        }
        info "添加设备控制: ID=$crid"
        add_device "$crid"
        show_firewall_status
        ;;
        
    "del")
        [ -z "$crid" ] && {
            echo "错误: 需要指定设备ID"
            exit 1
        }
        info "移除设备控制: ID=$crid"
        del_device "$crid"
        show_firewall_status
        ;;
        
    "status")
        show_firewall_status
        ;;
        
    "diagnose")
        diagnose
        ;;
        
    "test")
        # 测试地址解析
        echo "测试地址解析:"
        for test in "192.168.1.100" "192.168.1.0/24" "00:11:22:33:44:55" "invalid"; do
            echo -n "$test: "
            if parse_target "$test" >/dev/null; then
                echo "✓ 有效"
                parse_target "$test"
            else
                echo "✗ 无效"
            fi
        done
        ;;
        
    "flush")
        # 清理所有连接
        info "清理所有连接"
        if [ -x "$bin_conntrack" ]; then
            $bin_conntrack -F
            info "连接已清理"
        else
            info "conntrack不可用"
        fi
        ;;
        
    "help"|"")
        echo "时间控制系统命令工具"
        echo ""
        echo "用法: $0 {start|stop|add <id>|del <id>|status|diagnose|test|flush|help}"
        echo ""
        echo "命令说明:"
        echo "  start           - 初始化防火墙规则"
        echo "  stop            - 停止并清理所有防火墙规则"
        echo "  add <id>        - 添加设备到控制列表"
        echo "  del <id>        - 从控制列表移除设备"
        echo "  status          - 显示防火墙状态"
        echo "  diagnose        - 系统诊断"
        echo "  test            - 测试地址解析"
        echo "  flush           - 清理所有网络连接"
        echo "  help            - 显示此帮助信息"
        ;;
        
    *)
        echo "错误: 未知命令 '$crrun'"
        echo "使用: $0 help 查看帮助"
        exit 1
        ;;
esac