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 }