dns-server/threat/alert.go

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
}