增加威胁域名审计

threat/database.go

@@ -3,9 +3,13 @@ package threat
 import (
 	"encoding/csv"
 	"os"
+	"path/filepath"
 	"strconv"
 	"strings"
 	"sync"
+	"time"
+
+	"github.com/fsnotify/fsnotify"
 
 	"dns-server/logger"
 )
@@ -23,6 +27,10 @@ type ThreatDatabaseManager struct {
 	databasePath string
 	threatData   map[string]*ThreatInfo // 域名 -> 威胁信息
 	mutex        sync.RWMutex
+	watcher      *fsnotify.Watcher
+	watcherDone  chan struct{}
+	watcherMutex sync.Mutex
+	isWatching   bool
 }
 
 // NewThreatDatabaseManager 创建威胁域名数据库管理器
@@ -30,6 +38,7 @@ func NewThreatDatabaseManager(databasePath string) *ThreatDatabaseManager {
 	return &ThreatDatabaseManager{
 		databasePath: databasePath,
 		threatData:   make(map[string]*ThreatInfo),
+		watcherDone:  make(chan struct{}),
 	}
 }
 
@@ -192,3 +201,111 @@ func (m *ThreatDatabaseManager) saveDatabase() error {
 
 	return nil
 }
+
+// StartWatching 开始监听数据库文件变化
+func (m *ThreatDatabaseManager) StartWatching() error {
+	m.watcherMutex.Lock()
+	defer m.watcherMutex.Unlock()
+
+	if m.isWatching {
+		logger.Info("文件监听已经在运行中")
+		return nil
+	}
+
+	var err error
+	m.watcher, err = fsnotify.NewWatcher()
+	if err != nil {
+		logger.Error("创建文件监听器失败", "error", err)
+		return err
+	}
+
+	// 确保目录存在
+	dir := filepath.Dir(m.databasePath)
+	if _, err := os.Stat(dir); os.IsNotExist(err) {
+		logger.Error("数据库文件目录不存在", "dir", dir, "error", err)
+		return err
+	}
+
+	// 添加目录监听
+	if err := m.watcher.Add(dir); err != nil {
+		logger.Error("添加目录监听失败", "dir", dir, "error", err)
+		m.watcher.Close()
+		return err
+	}
+
+	m.isWatching = true
+	logger.Info("开始监听威胁域名数据库文件", "path", m.databasePath)
+
+	// 启动监听协程
+	go m.watchFileChanges()
+
+	return nil
+}
+
+// StopWatching 停止监听数据库文件变化
+func (m *ThreatDatabaseManager) StopWatching() {
+	m.watcherMutex.Lock()
+	defer m.watcherMutex.Unlock()
+
+	if !m.isWatching {
+		return
+	}
+
+	close(m.watcherDone)
+	if m.watcher != nil {
+		m.watcher.Close()
+	}
+	m.isWatching = false
+	logger.Info("停止监听威胁域名数据库文件")
+}
+
+// watchFileChanges 监听文件变化并重新加载数据库
+func (m *ThreatDatabaseManager) watchFileChanges() {
+	var reloadTimer *time.Timer
+	defer func() {
+		if reloadTimer != nil {
+			reloadTimer.Stop()
+		}
+	}()
+
+	for {
+		select {
+		case event, ok := <-m.watcher.Events:
+			if !ok {
+				return
+			}
+
+			// 检查是否是我们关注的文件
+			if filepath.Clean(event.Name) != filepath.Clean(m.databasePath) {
+				continue
+			}
+
+			logger.Debug("监听到文件变化", "event", event.Op.String(), "file", event.Name)
+
+			// 只关注文件写入和创建事件
+			if event.Has(fsnotify.Write) || event.Has(fsnotify.Create) {
+				// 使用防抖机制，避免频繁重新加载
+				if reloadTimer != nil {
+					reloadTimer.Stop()
+				}
+				reloadTimer = time.AfterFunc(500*time.Millisecond, func() {
+					logger.Info("检测到数据库文件变化，开始重新加载", "path", m.databasePath)
+					if err := m.LoadDatabase(); err != nil {
+						logger.Error("重新加载数据库失败", "error", err)
+					} else {
+						logger.Info("数据库重新加载成功")
+					}
+				})
+			}
+
+		case err, ok := <-m.watcher.Errors:
+			if !ok {
+				return
+			}
+			logger.Error("文件监听器错误", "error", err)
+
+		case <-m.watcherDone:
+			return
+		}
+	}
+}