dns-server/threat/database.go

package threat

import (
	"encoding/csv"
	"os"
	"strconv"
	"strings"
	"sync"

	"dns-server/logger"
)

// ThreatInfo 威胁域名详细信息
type ThreatInfo struct {
	Type       string // 威胁类型（如：钓鱼网站、木马、仿冒软件等）
	Name       string // 威胁名称（如：Silver fox 团伙）
	RiskLevel  int    // 风险等级（1-3，1=低，2=中，3=高）
	Domain     string // 域名
}

// ThreatDatabaseManager 威胁域名数据库管理器
type ThreatDatabaseManager struct {
	databasePath string
	threatData   map[string]*ThreatInfo // 域名 -> 威胁信息
	mutex        sync.RWMutex
}

// NewThreatDatabaseManager 创建威胁域名数据库管理器
func NewThreatDatabaseManager(databasePath string) *ThreatDatabaseManager {
	return &ThreatDatabaseManager{
		databasePath: databasePath,
		threatData:   make(map[string]*ThreatInfo),
	}
}

// LoadDatabase 加载威胁域名数据库
func (m *ThreatDatabaseManager) LoadDatabase() error {
	logger.Info("开始加载威胁域名数据库", "path", m.databasePath)

	// 打开CSV文件
	file, err := os.Open(m.databasePath)
	if err != nil {
		// 如果文件不存在，创建一个空文件
		if os.IsNotExist(err) {
			file, err = os.Create(m.databasePath)
			if err != nil {
				logger.Error("创建威胁域名数据库文件失败", "error", err)
				return err
			}
			file.Close()
			logger.Info("创建了空的威胁域名数据库文件")
			return nil
		}
		logger.Error("打开威胁域名数据库文件失败", "error", err)
		return err
	}
	defer file.Close()

	// 读取CSV文件
	reader := csv.NewReader(file)
	records, err := reader.ReadAll()
	if err != nil {
		logger.Error("读取威胁域名数据库文件失败", "error", err)
		return err
	}

	// 加载威胁域名
	m.mutex.Lock()
	defer m.mutex.Unlock()

	count := 0
	skipHeader := true
	for _, record := range records {
		// 跳过表头行
		if skipHeader {
			skipHeader = false
			continue
		}

		if len(record) >= 4 { // 确保至少有4列
			threatType := record[0] // 威胁类型
			name := record[1]       // 威胁名称
			riskLevelStr := record[2] // 风险等级
			domain := record[3]     // 域名

			// 解析风险等级
			riskLevel, err := strconv.Atoi(riskLevelStr)
			if err != nil {
				riskLevel = 1 // 默认低风险
			}

			// 存储威胁信息
			m.threatData[strings.ToLower(domain)] = &ThreatInfo{
				Type:      threatType,
				Name:      name,
				RiskLevel: riskLevel,
				Domain:    domain,
			}
			count++
		}
	}

	logger.Info("威胁域名数据库加载完成", "count", count)
	return nil
}

// GetThreatInfo 获取域名的威胁信息
func (m *ThreatDatabaseManager) GetThreatInfo(domain string) *ThreatInfo {
	m.mutex.RLock()
	defer m.mutex.RUnlock()
	return m.threatData[strings.ToLower(domain)]
}

// IsThreatDomain 检查域名是否在威胁数据库中
func (m *ThreatDatabaseManager) IsThreatDomain(domain string) bool {
	m.mutex.RLock()
	defer m.mutex.RUnlock()
	_, exists := m.threatData[strings.ToLower(domain)]
	return exists
}

// AddThreatDomain 添加威胁域名到数据库
func (m *ThreatDatabaseManager) AddThreatDomain(threatType, name string, riskLevel int, domain string) error {
	m.mutex.Lock()
	m.threatData[strings.ToLower(domain)] = &ThreatInfo{
		Type:      threatType,
		Name:      name,
		RiskLevel: riskLevel,
		Domain:    domain,
	}
	m.mutex.Unlock()

	// 保存到文件
	return m.saveDatabase()
}

// RemoveThreatDomain 从数据库中移除威胁域名
func (m *ThreatDatabaseManager) RemoveThreatDomain(domain string) error {
	m.mutex.Lock()
	delete(m.threatData, strings.ToLower(domain))
	m.mutex.Unlock()

	// 保存到文件
	return m.saveDatabase()
}

// GetAllThreatDomains 获取所有威胁域名信息
func (m *ThreatDatabaseManager) GetAllThreatDomains() []*ThreatInfo {
	m.mutex.RLock()
	defer m.mutex.RUnlock()

	threats := make([]*ThreatInfo, 0, len(m.threatData))
	for _, info := range m.threatData {
		threats = append(threats, info)
	}

	return threats
}

// saveDatabase 保存数据库到文件
func (m *ThreatDatabaseManager) saveDatabase() error {
	// 打开CSV文件
	file, err := os.Create(m.databasePath)
	if err != nil {
		return err
	}
	defer file.Close()

	// 写入CSV文件
	writer := csv.NewWriter(file)
	defer writer.Flush()

	// 写入表头
	if err := writer.Write([]string{"type", "name", "riskLevel", "domain"}); err != nil {
		return err
	}

	m.mutex.RLock()
	defer m.mutex.RUnlock()

	for _, info := range m.threatData {
		record := []string{
			info.Type,
			info.Name,
			strconv.Itoa(info.RiskLevel),
			info.Domain,
		}
		if err := writer.Write(record); err != nil {
			return err
		}
	}

	return nil
}