whois

threat/database.go

@@ -0,0 +1,194 @@
+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
+}