whois

threat/engine.go

@@ -0,0 +1,123 @@
+package threat
+
+import (
+	"dns-server/config"
+	"sync"
+	"time"
+
+	"dns-server/logger"
+)
+
+// ThreatEngine 威胁检测引擎
+type ThreatEngine struct {
+	config      *config.ThreatConfig
+	alertMgr    *AlertManager
+	dbManager   *ThreatDatabaseManager
+	clientStats map[string]*ClientQueryStats
+	mutex       sync.RWMutex
+}
+
+// NewThreatEngine 创建威胁检测引擎
+func NewThreatEngine(config *config.ThreatConfig, alertMgr *AlertManager, dbManager *ThreatDatabaseManager) *ThreatEngine {
+	return &ThreatEngine{
+		config:      config,
+		alertMgr:    alertMgr,
+		dbManager:   dbManager,
+		clientStats: make(map[string]*ClientQueryStats),
+	}
+}
+
+// CheckQuery 检查DNS查询是否存在威胁
+func (e *ThreatEngine) CheckQuery(sourceIP, domain, queryType string) []*ThreatAlert {
+	var alerts []*ThreatAlert
+
+	logger.Debug("威胁检测引擎检查域名", "domain", domain, "sourceIP", sourceIP)
+
+	// 检查威胁域名数据库
+	if alert := e.checkThreatDatabase(sourceIP, domain); alert != nil {
+		alerts = append(alerts, alert)
+	}
+
+	// 更新客户端查询统计
+	e.updateClientStats(sourceIP, queryType)
+
+	return alerts
+}
+
+// checkThreatDatabase 检查威胁域名数据库
+func (e *ThreatEngine) checkThreatDatabase(sourceIP, domain string) *ThreatAlert {
+	if e.dbManager == nil {
+		return nil
+	}
+
+	// 获取域名的威胁信息
+	threatInfo := e.dbManager.GetThreatInfo(domain)
+	if threatInfo == nil {
+		return nil
+	}
+
+	logger.Info("检测到威胁域名", "domain", domain, "type", threatInfo.Type, "name", threatInfo.Name, "riskLevel", threatInfo.RiskLevel)
+
+	// 根据风险等级确定告警级别
+	var alertLevel string
+	switch threatInfo.RiskLevel {
+	case 3:
+		alertLevel = AlertLevelHigh
+	case 2:
+		alertLevel = AlertLevelMedium
+	default:
+		alertLevel = AlertLevelLow
+	}
+
+	return &ThreatAlert{
+		ID:          generateAlertID(),
+		Timestamp:   time.Now(),
+		Level:       alertLevel,
+		Type:        AlertTypeSuspiciousDomain,
+		Description: "威胁域名数据库匹配",
+		Details:     "威胁类型: " + threatInfo.Type + ", 威胁名称: " + threatInfo.Name,
+		SourceIP:    sourceIP,
+		Domain:      domain,
+		QueryType:   "",
+		Resolved:    false,
+	}
+}
+
+// updateClientStats 更新客户端查询统计
+func (e *ThreatEngine) updateClientStats(sourceIP, queryType string) {
+	e.mutex.Lock()
+	defer e.mutex.Unlock()
+
+	stats, exists := e.clientStats[sourceIP]
+	if !exists {
+		stats = &ClientQueryStats{
+			QueryCount:    1,
+			NXDomainCount: 0,
+			LastQueryTime: time.Now(),
+			QueryTypes:    make(map[string]int),
+		}
+		e.clientStats[sourceIP] = stats
+	} else {
+		stats.QueryCount++
+		stats.LastQueryTime = time.Now()
+	}
+
+	// 更新查询类型统计
+	stats.QueryTypes[queryType]++
+}
+
+// generateAlertID 生成告警ID
+func generateAlertID() string {
+	return time.Now().Format("20060102150405") + "-" + randomString(8)
+}
+
+// randomString 生成随机字符串
+func randomString(length int) string {
+	const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+	result := make([]byte, length)
+	for i := range result {
+		result[i] = charset[time.Now().UnixNano()%int64(len(charset))]
+		time.Sleep(1 * time.Nanosecond)
+	}
+	return string(result)
+}