Initial commit
This commit is contained in:
kilyabin
@@ -0,0 +1,210 @@
|
||||
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
|
||||
Count int64
|
||||
}
|
||||
|
||||
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
|
||||
Count int64
|
||||
}
|
||||
|
||||
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
|
||||
Count int64
|
||||
}
|
||||
|
||||
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
|
||||
Count int64
|
||||
}
|
||||
|
||||
func (d *DB) AlertsTimeline(hours int) ([]TimelinePoint, error) {
|
||||
since := time.Now().Add(-time.Duration(hours) * time.Hour)
|
||||
var result []TimelinePoint
|
||||
err := d.db.Model(&Alert{}).
|
||||
Select("strftime('%Y-%m-%dT%H:00:00', created_at) as hour, count(*) as count").
|
||||
Where("created_at >= ?", since).
|
||||
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
|
||||
Count int64
|
||||
Bytes int64
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
package storage
|
||||
|
||||
import "time"
|
||||
|
||||
type Alert struct {
|
||||
ID uint `gorm:"primaryKey" json:"id"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
Type string `gorm:"index" json:"type"`
|
||||
Severity string `gorm:"index" json:"severity"`
|
||||
SrcIP string `gorm:"index" json:"src_ip"`
|
||||
DstIP string `json:"dst_ip"`
|
||||
SrcPort uint16 `json:"src_port"`
|
||||
DstPort uint16 `json:"dst_port"`
|
||||
Protocol string `json:"protocol"`
|
||||
Description string `json:"description"`
|
||||
Count int `json:"count"`
|
||||
}
|
||||
|
||||
type PacketStat struct {
|
||||
ID uint `gorm:"primaryKey" json:"id"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
Protocol string `json:"protocol"`
|
||||
SrcIP string `json:"src_ip"`
|
||||
DstIP string `json:"dst_ip"`
|
||||
Bytes int64 `json:"bytes"`
|
||||
}
|
||||
|
||||
type CaptureSession struct {
|
||||
ID uint `gorm:"primaryKey" json:"id"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
EndedAt *time.Time `json:"ended_at"`
|
||||
Source string `json:"source"`
|
||||
Mode string `json:"mode"`
|
||||
Status string `json:"status"`
|
||||
TotalPkts int64 `json:"total_pkts"`
|
||||
TotalBytes int64 `json:"total_bytes"`
|
||||
}
|
||||
|
||||
// AlertFilter используется для фильтрации алертов в запросах
|
||||
type AlertFilter struct {
|
||||
Type string
|
||||
Severity string
|
||||
SrcIP string
|
||||
DateFrom *time.Time
|
||||
DateTo *time.Time
|
||||
Page int
|
||||
PageSize int
|
||||
}
|
||||
|
||||
// SeverityOrder возвращает числовой приоритет для сортировки
|
||||
func SeverityOrder(s string) int {
|
||||
switch s {
|
||||
case "critical":
|
||||
return 4
|
||||
case "high":
|
||||
return 3
|
||||
case "medium":
|
||||
return 2
|
||||
case "low":
|
||||
return 1
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user