#!/bin/bash

# OpenWrt 防火牆規則查看腳本
# 用途：SSH 連線到防火牆並格式化顯示 nftables 規則

ROUTER_IP="192.168.42.10"
ROUTER_USER="root"  # 如果使用其他帳號請修改

# 顏色定義
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
GRAY='\033[0;90m'
NC='\033[0m' # No Color

echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
echo -e "${BLUE}🔒 OpenWrt 防火牆規則檢視工具${NC}"
echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
echo -e "${GRAY}目標設備: ${ROUTER_IP}${NC}"
echo -e "${GRAY}指令: nft list ruleset${NC}"
echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
echo ""

# 取得規則並暫存到 /tmp
TEMP_FILE=$(mktemp)
ssh ${ROUTER_USER}@${ROUTER_IP} "nft list ruleset" > "$TEMP_FILE" 2>&1

# 檢查 SSH 是否成功
if [ $? -ne 0 ]; then
    echo -e "${RED}❌ SSH 連線失敗！${NC}"
    echo -e "${YELLOW}請確認：${NC}"
    echo -e "  1. 防火牆設備是否開機 (${ROUTER_IP})"
    echo -e "  2. SSH 服務是否啟用"
    echo -e "  3. 網路連線是否正常"
    echo -e "  4. SSH 金鑰或密碼設定"
    rm -f "$TEMP_FILE"
    exit 1
fi

# 解析並顯示規則
current_section=""
in_table=false
in_chain=false
in_set=false
in_rule=false

while IFS= read -r line; do
    # 跳過空行（但保留段落分隔）
    if [[ -z "$line" ]]; then
        echo ""
        continue
    fi

    # Table 定義
    if [[ "$line" =~ ^table[[:space:]] ]]; then
        echo -e "\n${GREEN}╭─────────────────────────────────────────────────────────${NC}"
        echo -e "${GREEN}│ ${line}${NC}"
        echo -e "${GREEN}╰─────────────────────────────────────────────────────────${NC}"
        in_table=true
        continue
    fi

    # Chain 定義
    if [[ "$line" =~ ^[[:space:]]*chain[[:space:]] ]]; then
        chain_name=$(echo "$line" | sed -E 's/.*chain[[:space:]]+(\w+).*/\1/')
        echo -e "\n${YELLOW}  ▶ Chain: ${CYAN}${chain_name}${NC}"
        in_chain=true
        continue
    fi

    # Set 定義
    if [[ "$line" =~ ^[[:space:]]*set[[:space:]] ]]; then
        set_name=$(echo "$line" | sed -E 's/.*set[[:space:]]+(\w+).*/\1/')
        echo -e "\n${BLUE}  📦 Set: ${CYAN}${set_name}${NC}"
        # 顯示 set 的類型等資訊
        set_info=$(echo "$line" | sed -E 's/[[:space:]]*set[[:space:]]+(\{.*)/\1/')
        echo -e "${GRAY}     ${set_info}${NC}"
        in_set=true
        continue
    fi

    # 規則內容
    if [[ "$line" =~ ^[[:space:]]*meta[[:space:]] ]] || [[ "$line" =~ ^[[:space:]]*ip[[:space:]] ]] || [[ "$line" =~ ^[[:space:]]*tcp[[:space:]] ]] || [[ "$line" =~ ^[[:space:]]*udp[[:space:]] ]] || [[ "$line" =~ ^[[:space:]]*icmp[[:space:]] ]] || [[ "$line" =~ ^[[:space:]]*ct[[:space:]] ]]; then
        # 規則動作著色
        if [[ "$line" =~ accept ]]; then
            action_color="${GREEN}"
        elif [[ "$line" =~ drop|reject ]]; then
            action_color="${RED}"
        elif [[ "$line" =~ jump ]]; then
            action_color="${YELLOW}"
        else
            action_color="${NC}"
        fi

        # 著色顯示規則
        colored_line="$line"
        [[ "$line" =~ accept ]] && colored_line=$(echo "$line" | sed -E "s/accept/${GREEN}accept${NC}/g")
        [[ "$line" =~ drop ]] && colored_line=$(echo "$line" | sed -E "s/drop/${RED}drop${NC}/g")
        [[ "$line" =~ reject ]] && colored_line=$(echo "$line" | sed -E "s/reject/${RED}reject${NC}/g")
        [[ "$line" =~ jump ]] && colored_line=$(echo "$line" | sed -E "s/jump[[:space:]]+(\w+)/jump ${YELLOW}\1${NC}/")

        echo -e "    ${colored_line}"
        continue
    fi

    # Set 元素 (IP 位址)
    if [[ "$line" =~ ^[[:space:]]*[0-9] ]] || [[ "$line" =~ ^[[:space:]]*elements[[:space:]]*= ]]; then
        if [[ "$line" =~ elements[[:space:]]*= ]]; then
            echo -e "    ${GRAY}Elements:${NC}"
        else
            echo -e "    ${line}"
        fi
        continue
    fi

    # 其他設定
    if [[ "$line" =~ ^[[:space:]]*\}[[:space:]]*$ ]]; then
        # 結束符號，不顯示
        continue
    fi

    # policy 等
    if [[ "$line" =~ policy ]]; then
        policy=$(echo "$line" | sed -E 's/.*policy[[:space:]]+(\w+).*/\1/')
        if [[ "$policy" == "accept" ]]; then
            echo -e "    ${GRAY}policy: ${GREEN}${policy}${NC}"
        else
            echo -e "    ${GRAY}policy: ${RED}${policy}${NC}"
        fi
        continue
    fi

    # 其他內容直接顯示
    if [[ ! -z "$line" ]]; then
        echo -e "    ${GRAY}${line}${NC}"
    fi

done < "$TEMP_FILE"

rm -f "$TEMP_FILE"

echo -e "\n${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
echo -e "${GREEN}✓ 檢視完成${NC}"
echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"