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

START=99
STOP=10
USE_PROCD=1
EXTRA_COMMANDS="add_route"
EXTRA_HELP="	add_route	Add routing rules and policy"

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; }

    if [ "$proxy_server_ip_addr" = "127.0.0.1" ] || [ "$proxy_server_ip_addr" = "0.0.0.0" ]; then
        log_msg error "Invalid proxy_server_ip_addr: $proxy_server_ip_addr"; return 1;
    fi

    # 纵深防御：严禁将代理服务器设为本路由器自身的任一 IP，防止策略路由自指死循环导致网络崩溃
    if ip -4 -o addr show 2>/dev/null | awk '{print $4}' | cut -d/ -f1 | grep -Fxq "$proxy_server_ip_addr"; then
        log_msg error "proxy_server_ip_addr ($proxy_server_ip_addr) is a local IP of this router, aborting to prevent routing loop"
        return 1
    fi

    # 检查路由表 ID 是否为系统保留表 (0, 253, 254, 255)
    local routing_table=$(uci -q get "$UCI_CONFIG.global.routing_table" || echo "888")
    if [ "$routing_table" = "0" ] || [ "$routing_table" = "253" ] || [ "$routing_table" = "254" ] || [ "$routing_table" = "255" ]; then
        log_msg error "routing_table ($routing_table) is a reserved system table (0, 253, 254, 255), aborting"
        return 1
    fi

    return 0
}

probe_proxy_server() {
    local proxy_ip="$1"
    local proxy_dns_port="$2"
    local interface="$3"

    log_msg info "Probing proxy server $proxy_ip on interface $interface..."

    # 1. 检查 ARP 缓存或快速唤醒以获取 MAC 地址
    local mac=""
    if [ -f "/proc/net/arp" ]; then
        mac=$(awk -v ip="$proxy_ip" '$1 == ip { print $4; exit }' /proc/net/arp)
    fi

    if [ -n "$mac" ] && [ "$mac" != "00:00:00:00:00:00" ]; then
        log_msg info "Proxy server is REACHABLE: IP=$proxy_ip, MAC=$mac (found in ARP cache)"
    else
        # 快速单发 ping 1 秒唤醒 ARP 学习
        ping -c 1 -W 1 "$proxy_ip" >/dev/null 2>&1
        if [ -f "/proc/net/arp" ]; then
            mac=$(awk -v ip="$proxy_ip" '$1 == ip { print $4; exit }' /proc/net/arp)
        fi
        if [ -n "$mac" ] && [ "$mac" != "00:00:00:00:00:00" ]; then
            log_msg info "Proxy server is REACHABLE: IP=$proxy_ip, MAC=$mac"
        else
            log_msg warn "Proxy server $proxy_ip is NOT responding to ARP yet (it may be booting or offline, traffic will still be routed via policy table)"
        fi
    fi

    # 2. 若启用了上游 DNS 代理，在日志中记录目标地址与端口
    local dns_proxy=$(uci -q get "$UCI_CONFIG.global.dns_proxy_enabled" || echo 0)
    if [ "$dns_proxy" = "1" ]; then
        log_msg info "Upstream DNS redirection target configured: $proxy_ip#$proxy_dns_port"
    fi
}

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() {
    check_config || return 1

    local proxy_server_ip_addr=$(uci -q get "$UCI_CONFIG.global.proxy_server_ip_addr")
    local interface=$(uci -q get "$UCI_CONFIG.global.interface")
    interface="${interface:-br-lan}"
    local traffic_mark=$(uci -q get "$UCI_CONFIG.global.traffic_mark")
    traffic_mark="${traffic_mark:-0x666}"
    local routing_table=$(uci -q get "$UCI_CONFIG.global.routing_table")
    routing_table="${routing_table:-888}"
    local rule_priority=$(uci -q get "$UCI_CONFIG.global.rule_priority")
    rule_priority="${rule_priority:-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

    # 5. 禁用目标网卡的 ICMP 重定向通告，防止客户端绕行或 TCP 连接挂起
    if [ -d "/proc/sys/net/ipv4/conf/$interface" ]; then
        sysctl -w "net.ipv4.conf.$interface.send_redirects=0" >/dev/null 2>&1
        log_msg debug "Disabled send_redirects for $interface"
    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")
    curr_mark="${curr_mark:-0x666}"
    local curr_table=$(uci -q get "$UCI_CONFIG.global.routing_table")
    curr_table="${curr_table:-888}"
    local curr_prio=$(uci -q get "$UCI_CONFIG.global.rule_priority")
    curr_prio="${curr_prio:-1000}"
    local curr_iface=$(uci -q get "$UCI_CONFIG.global.interface")
    curr_iface="${curr_iface:-br-lan}"

    [ "$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
        if [ "$t" = "0" ] || [ "$t" = "253" ] || [ "$t" = "254" ] || [ "$t" = "255" ]; then
            log_msg error "Refusing to flush reserved system routing table $t"
            return
        fi
        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"

    # 恢复 ICMP 重定向通告
    if [ -d "/proc/sys/net/ipv4/conf/$curr_iface" ]; then
        sysctl -w "net.ipv4.conf.$curr_iface.send_redirects=1" >/dev/null 2>&1
    fi
}

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 via drop-in config..."

    mkdir -p /tmp/dnsmasq.d
    cat > /tmp/dnsmasq.d/flowproxy.conf << EOF
# Generated by FlowProxy - DO NOT EDIT MANUALLY
no-resolv
server=$proxy_server_ip_addr#$proxy_server_dns_port
EOF

    # 同步分发至各个可能存在的 dnsmasq 实例目录
    for dir in /tmp/dnsmasq*.d; do
        if [ -d "$dir" ] && [ "$dir" != "/tmp/dnsmasq.d" ]; then
            cp -f /tmp/dnsmasq.d/flowproxy.conf "$dir/flowproxy.conf"
        fi
    done

    # 历史遗留自愈：若早期版本存在残留的 dns_backup 节，静默移除避免脏数据
    if uci -q get "$UCI_CONFIG.dns_backup" >/dev/null; then
        uci -q delete "$UCI_CONFIG.dns_backup"
        uci -q commit "$UCI_CONFIG"
    fi

    /etc/init.d/dnsmasq restart
    log_msg info "DNS drop-in configuration applied successfully"
}

cleanup_dns() {
    log_msg info "Cleaning up DNS redirection drop-in configuration..."
    rm -f /tmp/dnsmasq.d/flowproxy.conf
    for dir in /tmp/dnsmasq*.d; do
        [ -d "$dir" ] && rm -f "$dir/flowproxy.conf"
    done

    # 历史遗留自愈
    if uci -q get "$UCI_CONFIG.dns_backup" >/dev/null; then
        uci -q delete "$UCI_CONFIG.dns_backup"
        uci -q commit "$UCI_CONFIG"
    fi

    /etc/init.d/dnsmasq restart
}

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

    local proxy_server_ip_addr=$(uci -q get "$UCI_CONFIG.global.proxy_server_ip_addr")
    local proxy_server_dns_port=$(uci -q get "$UCI_CONFIG.global.proxy_server_dns_port" || echo 5353)
    local interface=$(uci -q get "$UCI_CONFIG.global.interface" || echo "br-lan")
    probe_proxy_server "$proxy_server_ip_addr" "$proxy_server_dns_port" "$interface"

    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"
}

reload_service() {
    log_msg info "Reloading FlowProxy..."
    stop_service
    start_service
}

restart_service() {
    stop_service
    start_service
}

service_triggers() {
    # 监听自身配置改动
    procd_add_reload_trigger "$UCI_CONFIG"
    # 监听防火墙重载事件，防止规则被 fw4 清空后无法自动恢复
    procd_add_reload_trigger "firewall"
    # 监听网络接口事件
    local iface=$(uci -q get "$UCI_CONFIG.global.interface" || echo "br-lan")
    procd_add_reload_interface_trigger "$iface"
}