#!/bin/sh /etc/rc.common
# FlowProxy - Traffic Diversion Service

START=99
STOP=10
USE_PROCD=1

NAME="flowproxy"
UCI_CONFIG="flowproxy"
LOG_FILE="/var/log/flowproxy.log"
NFT_CONF="/tmp/flowproxy/nft.conf"
GENERATE_SCRIPT="/usr/share/flowproxy/generate_nft.sh"
STATE_FILE="/var/run/flowproxy.running"

get_level_weight() {
    case "$1" in
        debug) echo 0 ;;
        info)  echo 1 ;;
        warn)  echo 2 ;;
        error) echo 3 ;;
        *)     echo 1 ;;
    esac
}

log_msg() {
    local level="$1"; shift
    local msg="$*"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    
    # 等级过滤
    local conf_level=$(uci -q get "$UCI_CONFIG.global.log_level" || echo "info")
    local conf_weight=$(get_level_weight "$conf_level")
    local msg_weight=$(get_level_weight "$level")
    [ "$msg_weight" -lt "$conf_weight" ] && return

    # 简单的日志截断逻辑，防止日志无限增长
    # 读取配置的大小限制，默认 1024KB
    local max_size=$(uci -q get "$UCI_CONFIG.global.log_size" || echo 1024)
    local current_size=$(du -k "$LOG_FILE" 2>/dev/null | cut -f1)
    if [ "${current_size:-0}" -gt "$max_size" ]; then
        mv "$LOG_FILE" "$LOG_FILE.old"
        touch "$LOG_FILE"
    fi

    echo "[$timestamp] [$level] $msg" >> "$LOG_FILE"
}

check_env() {
    # 检测核心二进制工具是否存在
    if ! command -v nft >/dev/null 2>&1; then
        log_msg error "nftables binary not found"
        return 1
    fi
    return 0
}

check_config() {
    local proxy_server_ip_addr=$(uci -q get "$UCI_CONFIG.global.proxy_server_ip_addr")
    [ -z "$proxy_server_ip_addr" ] && { log_msg error "proxy_server_ip_addr not configured"; return 1; }
    return 0
}

apply_nft_rules() {
    log_msg info "Generating nftables configuration..."
    local conf_level=$(uci -q get "$UCI_CONFIG.global.log_level" || echo "info")

    mkdir -p "$(dirname "$NFT_CONF")"
    if [ -x "$GENERATE_SCRIPT" ]; then
        # 传递日志等级并捕获 stderr 到临时文件
        local err_file="/tmp/flowproxy/gen_err.tmp"
        FLOWPROXY_LOG_LEVEL="$conf_level" "$GENERATE_SCRIPT" >/dev/null 2> "$err_file"
        local ret=$?
        
        # 将脚本的输出记录到日志
        if [ -s "$err_file" ]; then
            while read -r line; do
                log_msg debug "Script: $line"
            done < "$err_file"
            rm -f "$err_file"
        fi

        if [ $ret -eq 0 ]; then
            log_msg info "Applying nftables rules from $NFT_CONF"
            # 确保内核支持 inet 表
            nft delete table inet flowproxy 2>/dev/null
            if nft -f "$NFT_CONF"; then
                return 0
            else
                log_msg error "nftables rules application failed (syntax error)"
                return 1
            fi
        else
            log_msg error "Failed to generate nftables config"
            return 1
        fi
    else
        log_msg error "Generator script $GENERATE_SCRIPT not found or not executable"
        return 1
    fi
}

add_route() {
    local proxy_server_ip_addr=$(uci -q get "$UCI_CONFIG.global.proxy_server_ip_addr")
    local interface=$(uci -q get "$UCI_CONFIG.global.interface" || echo "br-lan")
    local traffic_mark=$(uci -q get "$UCI_CONFIG.global.traffic_mark" || echo "0x666")
    local routing_table=$(uci -q get "$UCI_CONFIG.global.routing_table" || echo "888")
    local rule_priority=$(uci -q get "$UCI_CONFIG.global.rule_priority" || echo "1000")

    # 1. 先尝试用旧记录清理一次，防止由于配置变更导致的残留
    delete_route "quiet"

    log_msg info "Setting up routing table $routing_table and rule for mark $traffic_mark with priority $rule_priority..."
    
    # 2. 记录当前使用的标记值、路由表 ID 和优先级
    echo "$traffic_mark:$routing_table:$rule_priority" > "/var/run/flowproxy.mark"

    # 3. 再次确保新标记和表的规则是干净的
    while ip rule del fwmark "$traffic_mark" lookup "$routing_table" 2>/dev/null; do :; done
    ip route flush table "$routing_table" 2>/dev/null

    # 4. 添加路由和策略
    local err
    ip route add default via "$proxy_server_ip_addr" dev "$interface" table "$routing_table" 2>/dev/null
    err=$(ip rule add fwmark "$traffic_mark" lookup "$routing_table" priority "$rule_priority" 2>&1)
    if [ $? -eq 0 ]; then
        log_msg info "Route added: via $proxy_server_ip_addr dev $interface table $routing_table for mark $traffic_mark with priority $rule_priority"
    else
        log_msg error "Failed to add ip rule: $err"
    fi
}

delete_route() {
    local traffic_mark routing_table rule_priority
    local quiet="$1"
    
    # 1. 获取两组清理目标：旧值(记录文件) 和 新值(当前UCI)
    local old_mark="" old_table="" old_prio=""
    if [ -f "/var/run/flowproxy.mark" ]; then
        local saved_val=$(cat "/var/run/flowproxy.mark")
        old_mark=$(echo "$saved_val" | cut -d':' -f1)
        old_table=$(echo "$saved_val" | cut -d':' -f2)
        old_prio=$(echo "$saved_val" | cut -d':' -f3)
        [ "$quiet" != "quiet" ] && rm -f "/var/run/flowproxy.mark"
    fi

    local curr_mark=$(uci -q get "$UCI_CONFIG.global.traffic_mark" || echo "0x666")
    local curr_table=$(uci -q get "$UCI_CONFIG.global.routing_table" || echo "888")
    local curr_prio=$(uci -q get "$UCI_CONFIG.global.rule_priority" || echo "1000")

    [ "$quiet" != "quiet" ] && log_msg info "Cleaning up FlowProxy routes (Atomic Mode)..."
    
    # 2. 核心清理函数：只清理带 fwmark 且指向特定 table 的规则
    safe_del_rule() {
        local m="$1" t="$2" p="$3"
        [ -z "$m" ] || [ -z "$t" ] && return
        log_msg debug "Attempting cleanup: mark=$m, table=$t, priority=${p:-any}"
        # 精准清理 (带优先级)
        [ -n "$p" ] && while ip rule del fwmark "$m" lookup "$t" priority "$p" 2>/dev/null; do 
            log_msg debug "Removed stale ip rule: mark $m lookup $t priority $p"
        done
        # 模糊清理 (不带优先级，防止配置中途被手动改动)
        while ip rule del fwmark "$m" lookup "$t" 2>/dev/null; do 
            log_msg debug "Removed stale ip rule (fuzzy): mark $m lookup $t"
        done
        # 清空路由表
        ip route flush table "$t" 2>/dev/null
        log_msg debug "Flushed routing table $t"
    }

    # 执行清理：先清旧值，再清当前值 (防止残留)
    safe_del_rule "$old_mark" "$old_table" "$old_prio"
    [ "$old_mark" != "$curr_mark" ] || [ "$old_table" != "$curr_table" ] && \
        safe_del_rule "$curr_mark" "$curr_table" "$curr_prio"
}

delete_nft() {
    nft delete table inet flowproxy 2>/dev/null
    log_msg info "Removed nftables table 'flowproxy'"
}

setup_dns() {
    local enabled=$(uci -q get "$UCI_CONFIG.global.enabled" || echo 0)
    local dns_proxy_enabled=$(uci -q get "$UCI_CONFIG.global.dns_proxy_enabled" || echo 0)
    local proxy_server_dns_port=$(uci -q get "$UCI_CONFIG.global.proxy_server_dns_port" || echo 5353)
    local proxy_server_ip_addr=$(uci -q get "$UCI_CONFIG.global.proxy_server_ip_addr")

    # Only apply DNS proxy if BOTH main service and DNS proxy are enabled
    if [ "$enabled" -eq 0 ] || [ "$dns_proxy_enabled" -eq 0 ]; then
        cleanup_dns
        return 0
    fi

    [ -z "$proxy_server_ip_addr" ] && { log_msg error "proxy_server_ip_addr not configured, cannot setup DNS"; return 1; }

    log_msg info "Setting up DNS redirection to $proxy_server_ip_addr#$proxy_server_dns_port..."

    # Persistent backup in UCI
    if ! uci -q get "$UCI_CONFIG.dns_backup" >/dev/null; then
        local old_noresolv=$(uci -q get dhcp.@dnsmasq[0].noresolv)
        local old_dns_redirect=$(uci -q get dhcp.@dnsmasq[0].dns_redirect)
        local old_servers=$(uci -q get dhcp.@dnsmasq[0].server | tr '\n' ',' | sed 's/,$//')
        
        log_msg debug "Backing up DNS: noresolv=${old_noresolv:-empty}, redirect=${old_dns_redirect:-empty}, servers=${old_servers:-empty}"
        
        uci set "$UCI_CONFIG.dns_backup=dns_backup"
        [ -n "$old_noresolv" ] && uci set "$UCI_CONFIG.dns_backup.noresolv=$old_noresolv"
        [ -n "$old_dns_redirect" ] && uci set "$UCI_CONFIG.dns_backup.dns_redirect=$old_dns_redirect"
        [ -n "$old_servers" ] && uci set "$UCI_CONFIG.dns_backup.servers=$old_servers"
        uci commit "$UCI_CONFIG"
    fi

    log_msg debug "Applying new DNS: noresolv=1, redirect=0, server=$proxy_server_ip_addr#$proxy_server_dns_port"
    uci set dhcp.@dnsmasq[0].noresolv='1'
    uci set dhcp.@dnsmasq[0].dns_redirect='0'
    
    # Delete current servers and add our proxy
    uci -q delete dhcp.@dnsmasq[0].server
    uci add_list dhcp.@dnsmasq[0].server="$proxy_server_ip_addr#$proxy_server_dns_port"
    
    uci commit dhcp
    /etc/init.d/dnsmasq restart
}

cleanup_dns() {
    if uci -q get "$UCI_CONFIG.dns_backup" >/dev/null; then
        log_msg info "Restoring original DNS configuration..."
        
        local noresolv=$(uci -q get "$UCI_CONFIG.dns_backup.noresolv")
        local dns_redirect=$(uci -q get "$UCI_CONFIG.dns_backup.dns_redirect")
        local servers=$(uci -q get "$UCI_CONFIG.dns_backup.servers")
        
        log_msg debug "Restoring DNS values: noresolv=${noresolv:-0}, redirect=${dns_redirect:-1}, servers=${servers:-empty}"

        # Restore noresolv (default to 0 if not backed up)
        [ -n "$noresolv" ] && uci set dhcp.@dnsmasq[0].noresolv="$noresolv" || uci set dhcp.@dnsmasq[0].noresolv="0"

        # Restore dns_redirect (only restore if it originally existed, avoid polluting config)
        if [ -n "$dns_redirect" ]; then
            uci set dhcp.@dnsmasq[0].dns_redirect="$dns_redirect"
        else
            uci -q delete dhcp.@dnsmasq[0].dns_redirect
        fi
        
        # Restore servers correctly splitting comma-separated list
        uci -q delete dhcp.@dnsmasq[0].server
        if [ -n "$servers" ]; then
            for s in $(echo "$servers" | tr ',' ' '); do
                [ -n "$s" ] && uci add_list dhcp.@dnsmasq[0].server="$s"
            done
        fi
        
        uci commit dhcp
        uci delete "$UCI_CONFIG.dns_backup"
        uci commit "$UCI_CONFIG"
        /etc/init.d/dnsmasq restart
    fi
}

start_service() {
    local enabled=$(uci -q get "$UCI_CONFIG.global.enabled" || echo 0)

    if [ "$enabled" -eq 0 ]; then
        log_msg info "FlowProxy is disabled, ensuring service is stopped"
        stop_service
        return 0
    fi

    log_msg info "Starting FlowProxy..."

    check_env || return 1
    check_config || return 1
    apply_nft_rules || return 1
    add_route
    setup_dns

    touch "$STATE_FILE"
    log_msg info "FlowProxy service started successfully"
}
stop_service() {
    log_msg info "Stopping FlowProxy..."
    # 强制先移除标记，防止状态误报
    rm -f "$STATE_FILE"
    delete_route
    delete_nft
    cleanup_dns
    log_msg info "FlowProxy service stopped"
}

restart_service() {
    stop_service
    start_service
}

reload_service() {
    restart_service
}

service_triggers() {
    # 监听自身配置改动
    procd_add_reload_trigger "$UCI_CONFIG"
    # 监听防火墙重载事件，防止规则被 fw4 清空后无法自动恢复
    procd_add_reload_trigger "firewall"
}
