dip-ids/internal/storage/db.go

package storage

import (
	"fmt"
	"time"

	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

type DB struct {
	db *gorm.DB
}

func New(path string) (*DB, error) {
	db, err := gorm.Open(sqlite.Open(path), &gorm.Config{
		Logger: logger.Default.LogMode(logger.Silent),
	})
	if err != nil {
		return nil, fmt.Errorf("open database: %w", err)
	}

	if err := db.AutoMigrate(&Alert{}, &PacketStat{}, &CaptureSession{}); err != nil {
		return nil, fmt.Errorf("migrate: %w", err)
	}

	return &DB{db: db}, nil
}

// ── Alerts ──────────────────────────────────────────────────────────────────

func (d *DB) CreateAlert(a *Alert) error {
	return d.db.Create(a).Error
}

func (d *DB) GetAlerts(f AlertFilter) ([]Alert, int64, error) {
	q := d.db.Model(&Alert{})

	if f.Type != "" {
		q = q.Where("type = ?", f.Type)
	}
	if f.Severity != "" {
		q = q.Where("severity = ?", f.Severity)
	}
	if f.SrcIP != "" {
		q = q.Where("src_ip = ?", f.SrcIP)
	}
	if f.DateFrom != nil {
		q = q.Where("created_at >= ?", f.DateFrom)
	}
	if f.DateTo != nil {
		q = q.Where("created_at <= ?", f.DateTo)
	}

	var total int64
	q.Count(&total)

	if f.PageSize <= 0 {
		f.PageSize = 50
	}
	if f.Page <= 0 {
		f.Page = 1
	}

	var alerts []Alert
	err := q.Order("created_at desc").
		Offset((f.Page - 1) * f.PageSize).
		Limit(f.PageSize).
		Find(&alerts).Error

	return alerts, total, err
}

func (d *DB) DeleteAlert(id uint) error {
	return d.db.Delete(&Alert{}, id).Error
}

func (d *DB) DeleteAllAlerts() error {
	return d.db.Where("1 = 1").Delete(&Alert{}).Error
}

// ── Stats ────────────────────────────────────────────────────────────────────

func (d *DB) CreatePacketStat(s *PacketStat) error {
	return d.db.Create(s).Error
}

type SeverityCount struct {
	Severity string `json:"severity"`
	Count    int64  `json:"count"`
}

func (d *DB) AlertsBySeverity() ([]SeverityCount, error) {
	var result []SeverityCount
	err := d.db.Model(&Alert{}).
		Select("severity, count(*) as count").
		Group("severity").
		Find(&result).Error
	return result, err
}

type TypeCount struct {
	Type  string `json:"type"`
	Count int64  `json:"count"`
}

func (d *DB) AlertsByType() ([]TypeCount, error) {
	var result []TypeCount
	err := d.db.Model(&Alert{}).
		Select("type, count(*) as count").
		Group("type").
		Find(&result).Error
	return result, err
}

type IPCount struct {
	SrcIP string `json:"src_ip"`
	Count int64  `json:"count"`
}

func (d *DB) TopSourceIPs(limit int) ([]IPCount, error) {
	var result []IPCount
	err := d.db.Model(&Alert{}).
		Select("src_ip, count(*) as count").
		Group("src_ip").
		Order("count desc").
		Limit(limit).
		Find(&result).Error
	return result, err
}

type TimelinePoint struct {
	Hour  string `json:"hour"`
	Count int64  `json:"count"`
}

func (d *DB) AlertsTimeline(hours int) ([]TimelinePoint, error) {
	since := time.Now().UTC().Add(-time.Duration(hours) * time.Hour)
	var result []TimelinePoint
	err := d.db.Model(&Alert{}).
		Select("strftime('%Y-%m-%dT%H:00:00Z', created_at) as hour, count(*) as count").
		Where("created_at >= datetime(?)", since.Format("2006-01-02 15:04:05")).
		Group("hour").
		Order("hour asc").
		Find(&result).Error
	return result, err
}

func (d *DB) TotalAlerts() (int64, error) {
	var count int64
	err := d.db.Model(&Alert{}).Count(&count).Error
	return count, err
}

type ProtocolCount struct {
	Protocol string `json:"protocol"`
	Count    int64  `json:"count"`
	Bytes    int64  `json:"bytes"`
}

func (d *DB) TrafficByProtocol() ([]ProtocolCount, error) {
	var result []ProtocolCount
	err := d.db.Model(&PacketStat{}).
		Select("protocol, count(*) as count, sum(bytes) as bytes").
		Group("protocol").
		Order("count desc").
		Find(&result).Error
	return result, err
}

func (d *DB) GetRecentPacketStats(limit int) ([]PacketStat, error) {
	var stats []PacketStat
	err := d.db.Order("created_at desc").Limit(limit).Find(&stats).Error
	return stats, err
}

// ── Sessions ─────────────────────────────────────────────────────────────────

func (d *DB) CreateSession(s *CaptureSession) error {
	return d.db.Create(s).Error
}

func (d *DB) UpdateSession(s *CaptureSession) error {
	return d.db.Save(s).Error
}

func (d *DB) GetActiveSession() (*CaptureSession, error) {
	var s CaptureSession
	err := d.db.Where("status = ?", "running").First(&s).Error
	if err != nil {
		return nil, err
	}
	return &s, nil
}

func (d *DB) GetLastSession() (*CaptureSession, error) {
	var s CaptureSession
	err := d.db.Order("created_at desc").First(&s).Error
	if err != nil {
		return nil, err
	}
	return &s, nil
}

func (d *DB) GetAllSessions() ([]CaptureSession, error) {
	var sessions []CaptureSession
	err := d.db.Order("created_at desc").Limit(20).Find(&sessions).Error
	return sessions, err
}