dip-ids/internal/api/handlers/capture.go

package handlers

import (
	"context"
	"log"
	"mime/multipart"
	"net/http"
	"os"
	"path/filepath"
	"time"

	"dip-ids/internal/analyzer"
	"dip-ids/internal/capture"
	"dip-ids/internal/storage"

	"github.com/gin-gonic/gin"
)

// CaptureState хранит состояние текущей сессии захвата
type CaptureState struct {
	cancel    context.CancelFunc
	session   *storage.CaptureSession
	engine    *analyzer.Engine
}

type CaptureHandler struct {
	db          *storage.DB
	engine      *analyzer.Engine
	defaultIface string
	active      *CaptureState
}

func NewCaptureHandler(db *storage.DB, engine *analyzer.Engine, defaultIface string) *CaptureHandler {
	return &CaptureHandler{
		db:           db,
		engine:       engine,
		defaultIface: defaultIface,
	}
}

// StartLive POST /api/capture/start
func (h *CaptureHandler) StartLive(c *gin.Context) {
	if h.active != nil {
		c.JSON(http.StatusConflict, gin.H{"error": "capture already running"})
		return
	}

	var req struct {
		Interface string `json:"interface"`
	}
	_ = c.ShouldBindJSON(&req)
	if req.Interface == "" {
		req.Interface = h.defaultIface
	}

	sess := &storage.CaptureSession{
		CreatedAt: time.Now(),
		Source:    req.Interface,
		Mode:      "live",
		Status:    "running",
	}
	if err := h.db.CreateSession(sess); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	ctx, cancel := context.WithCancel(context.Background())
	cap, err := capture.New(capture.Config{
		Interface:   req.Interface,
		Mode:        capture.ModeLive,
		Snaplen:     65535,
		Promiscuous: true,
	})
	if err != nil {
		cancel()
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	pkts, err := cap.Start(ctx)
	if err != nil {
		cancel()
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	h.engine.Reset()
	h.active = &CaptureState{cancel: cancel, session: sess, engine: h.engine}

	go func() {
		h.engine.Run(ctx, pkts)
		// Обновляем сессию после завершения
		now := time.Now()
		pktsCount, bytesCount := h.engine.Stats()
		sess.EndedAt = &now
		sess.Status = "stopped"
		sess.TotalPkts = pktsCount
		sess.TotalBytes = bytesCount
		if err := h.db.UpdateSession(sess); err != nil {
			log.Printf("[capture] update session: %v", err)
		}
		h.active = nil
	}()

	c.JSON(http.StatusOK, gin.H{"session_id": sess.ID, "interface": req.Interface})
}

// Stop POST /api/capture/stop
func (h *CaptureHandler) Stop(c *gin.Context) {
	if h.active == nil {
		c.JSON(http.StatusConflict, gin.H{"error": "no active capture"})
		return
	}
	h.active.cancel()
	c.JSON(http.StatusOK, gin.H{"ok": true})
}

// UploadPCAP POST /api/capture/upload
func (h *CaptureHandler) UploadPCAP(c *gin.Context) {
	if h.active != nil {
		c.JSON(http.StatusConflict, gin.H{"error": "capture already running"})
		return
	}

	file, err := c.FormFile("file")
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "file required"})
		return
	}

	tmpPath, err := saveTempFile(file)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	sess := &storage.CaptureSession{
		CreatedAt: time.Now(),
		Source:    file.Filename,
		Mode:      "pcap",
		Status:    "running",
	}
	if err := h.db.CreateSession(sess); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	ctx, cancel := context.WithCancel(context.Background())
	cap, err := capture.New(capture.Config{
		PCAPFile: tmpPath,
		Mode:     capture.ModePCAP,
		Snaplen:  65535,
	})
	if err != nil {
		cancel()
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	pkts, err := cap.Start(ctx)
	if err != nil {
		cancel()
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	h.engine.Reset()
	h.active = &CaptureState{cancel: cancel, session: sess, engine: h.engine}

	go func() {
		defer os.Remove(tmpPath)
		h.engine.Run(ctx, pkts)

		now := time.Now()
		pktsCount, bytesCount := h.engine.Stats()
		sess.EndedAt = &now
		sess.Status = "completed"
		sess.TotalPkts = pktsCount
		sess.TotalBytes = bytesCount
		if err := h.db.UpdateSession(sess); err != nil {
			log.Printf("[capture] update session: %v", err)
		}
		h.active = nil
	}()

	c.JSON(http.StatusOK, gin.H{"session_id": sess.ID, "filename": file.Filename})
}

// Status GET /api/capture/status
func (h *CaptureHandler) Status(c *gin.Context) {
	if h.active == nil {
		last, _ := h.db.GetLastSession()
		c.JSON(http.StatusOK, gin.H{"running": false, "last_session": last})
		return
	}

	pkts, bytes := h.engine.Stats()
	c.JSON(http.StatusOK, gin.H{
		"running":    true,
		"session_id": h.active.session.ID,
		"source":     h.active.session.Source,
		"mode":       h.active.session.Mode,
		"pkts":       pkts,
		"bytes":      bytes,
	})
}

// GetInterfaces GET /api/capture/interfaces
func (h *CaptureHandler) GetInterfaces(c *gin.Context) {
	ifaces, err := capture.ListInterfaces()
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"interfaces": ifaces})
}

// GetSessions GET /api/capture/sessions
func (h *CaptureHandler) GetSessions(c *gin.Context) {
	sessions, err := h.db.GetAllSessions()
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"data": sessions})
}

func saveTempFile(fh *multipart.FileHeader) (string, error) {
	src, err := fh.Open()
	if err != nil {
		return "", err
	}
	defer src.Close()

	tmp, err := os.CreateTemp("", "nids-*.pcap")
	if err != nil {
		return "", err
	}
	defer tmp.Close()

	buf := make([]byte, 32*1024)
	for {
		n, err := src.Read(buf)
		if n > 0 {
			tmp.Write(buf[:n])
		}
		if err != nil {
			break
		}
	}

	return filepath.Abs(tmp.Name())
}