whois

threat/alert.go

@@ -0,0 +1,193 @@
+package threat
+
+import (
+	"dns-server/config"
+	"encoding/json"
+	"os"
+	"path/filepath"
+	"sort"
+	"sync"
+	"time"
+)
+
+// AlertManager 告警管理器
+type AlertManager struct {
+	alerts      []*ThreatAlert
+	mutex       sync.RWMutex
+	config      *config.ThreatConfig
+	storagePath string
+}
+
+// NewAlertManager 创建告警管理器
+func NewAlertManager(config *config.ThreatConfig) *AlertManager {
+	storagePath := "./data/threat_alerts.json"
+	
+	// 确保数据目录存在
+	dataDir := filepath.Dir(storagePath)
+	if err := os.MkdirAll(dataDir, 0755); err != nil {
+		// 记录错误但继续执行
+	}
+
+	return &AlertManager{
+		alerts:      make([]*ThreatAlert, 0),
+		config:      config,
+		storagePath: storagePath,
+	}
+}
+
+// AddAlert 添加告警
+func (m *AlertManager) AddAlert(alert *ThreatAlert) {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+
+	// 添加告警到内存
+	m.alerts = append(m.alerts, alert)
+
+	// 清理过期告警
+	m.cleanupExpiredAlerts()
+
+	// 持久化存储
+	m.saveAlerts()
+}
+
+// GetAlerts 获取告警列表
+func (m *AlertManager) GetAlerts(limit, offset int, level string) []*ThreatAlert {
+	m.mutex.RLock()
+	defer m.mutex.RUnlock()
+
+	// 过滤告警
+	var filteredAlerts []*ThreatAlert
+	for _, alert := range m.alerts {
+		if level == "" || alert.Level == level {
+			filteredAlerts = append(filteredAlerts, alert)
+		}
+	}
+
+	// 按时间排序（最新的在前）
+	sort.Slice(filteredAlerts, func(i, j int) bool {
+		return filteredAlerts[i].Timestamp.After(filteredAlerts[j].Timestamp)
+	})
+
+	// 应用分页
+	start := offset
+	end := offset + limit
+	if start < 0 {
+		start = 0
+	}
+	if end > len(filteredAlerts) {
+		end = len(filteredAlerts)
+	}
+	if start >= len(filteredAlerts) {
+		return []*ThreatAlert{}
+	}
+
+	return filteredAlerts[start:end]
+}
+
+// GetAlert 根据ID获取告警
+func (m *AlertManager) GetAlert(alertID string) *ThreatAlert {
+	m.mutex.RLock()
+	defer m.mutex.RUnlock()
+
+	for _, alert := range m.alerts {
+		if alert.ID == alertID {
+			return alert
+		}
+	}
+
+	return nil
+}
+
+// ResolveAlert 解决告警
+func (m *AlertManager) ResolveAlert(alertID, action string) bool {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+
+	for _, alert := range m.alerts {
+		if alert.ID == alertID {
+			alert.Resolved = true
+			alert.ResolvedTime = time.Now()
+			alert.Action = action
+			
+			// 持久化存储
+			m.saveAlerts()
+			return true
+		}
+	}
+
+	return false
+}
+
+// GetAlertCount 获取告警数量
+func (m *AlertManager) GetAlertCount(level string) int {
+	m.mutex.RLock()
+	defer m.mutex.RUnlock()
+
+	count := 0
+	for _, alert := range m.alerts {
+		if !alert.Resolved && (level == "" || alert.Level == level) {
+			count++
+		}
+	}
+
+	return count
+}
+
+// LoadAlerts 加载告警
+func (m *AlertManager) LoadAlerts() error {
+	// 检查文件是否存在
+	if _, err := os.Stat(m.storagePath); os.IsNotExist(err) {
+		return nil
+	}
+
+	// 读取文件
+	data, err := os.ReadFile(m.storagePath)
+	if err != nil {
+		return err
+	}
+
+	// 解析JSON
+	var alerts []*ThreatAlert
+	if err := json.Unmarshal(data, &alerts); err != nil {
+		return err
+	}
+
+	// 更新内存中的告警
+	m.mutex.Lock()
+	m.alerts = alerts
+	m.mutex.Unlock()
+
+	// 清理过期告警
+	m.cleanupExpiredAlerts()
+
+	return nil
+}
+
+// saveAlerts 保存告警
+func (m *AlertManager) saveAlerts() error {
+	data, err := json.MarshalIndent(m.alerts, "", "  ")
+	if err != nil {
+		return err
+	}
+
+	return os.WriteFile(m.storagePath, data, 0644)
+}
+
+// cleanupExpiredAlerts 清理过期告警
+func (m *AlertManager) cleanupExpiredAlerts() {
+	if m.config.AlertRetentionDays <= 0 {
+		return // 0表示一直保存
+	}
+
+	now := time.Now()
+	retentionTime := time.Duration(m.config.AlertRetentionDays) * 24 * time.Hour
+
+	var validAlerts []*ThreatAlert
+	for _, alert := range m.alerts {
+		if now.Sub(alert.Timestamp) <= retentionTime {
+			validAlerts = append(validAlerts, alert)
+		}
+	}
+
+	m.alerts = validAlerts
+}