go_gameshield/internal/repository/admin/waflog.go

package admin

import (
	"context"
	v1 "github.com/go-nunu/nunu-layout-advanced/api/v1"
	adminApi "github.com/go-nunu/nunu-layout-advanced/api/v1/admin"
	"github.com/go-nunu/nunu-layout-advanced/internal/model"
	"github.com/go-nunu/nunu-layout-advanced/internal/repository"
	"gorm.io/gorm"
	"math"
	"strings"
)

type WafLogRepository interface {
	GetWafLog(ctx context.Context, id int64) (*model.WafLog, error)
	GetWafLogList(ctx context.Context, req adminApi.SearchWafLogParams) (*v1.PaginatedResponse[model.WafLog], error)
	AddWafLog(ctx context.Context, log *model.WafLog) error
	BatchAddWafLog(ctx context.Context, logs []*model.WafLog) error
	ExportWafLog(ctx context.Context, req adminApi.ExportWafLog) ([]model.WafLogWithGatewayIP, error)
	ExportWafLogWithPagination(ctx context.Context, req adminApi.ExportWafLog, page, pageSize int) ([]model.WafLogWithGatewayIP, error)
	GetWafLogExportCount(ctx context.Context, req adminApi.ExportWafLog) (int, error)
}

func NewWafLogRepository(
	repository *repository.Repository,
) WafLogRepository {
	return &wafLogRepository{
		Repository: repository,
	}
}

type wafLogRepository struct {
	*repository.Repository
}

// buildExportQuery 是一个辅助函数，用于构建导出日志的公共查询条件
func (r *wafLogRepository) buildExportQuery(ctx context.Context, req adminApi.ExportWafLog) *gorm.DB {
	// 使用 Table("waf_log as wl") 是为了给主表起一个别名，方便子查询中引用
	query := r.DBWithName(ctx, "admin").Model(&model.WafLog{}).Table("waf_log as wl")

	if req.RequestIp != "" {
		query = query.Where("wl.request_ip = ?", strings.TrimSpace(req.RequestIp))
	}
	if req.Uid != 0 {
		query = query.Where("wl.uid = ?", req.Uid)
	}
	if req.Api != "" {
		query = query.Where("wl.api = ?", strings.TrimSpace(req.Api))
	}
	if req.Name != "" {
		query = query.Where("wl.name = ?", strings.TrimSpace(req.Name))
	}
	if req.RuleId != 0 {
		query = query.Where("wl.rule_id = ?", req.RuleId)
	}
	if len(req.HostIds) > 0 {
		query = query.Where("wl.host_id IN ?", req.HostIds)
	}
	if req.UserAgent != "" {
		query = query.Where("wl.user_agent = ?", strings.TrimSpace(req.UserAgent))
	}
	if len(req.ApiNames) > 0 {
		query = query.Where("wl.api_name IN ?", req.ApiNames)
	}
	if len(req.ApiTypes) > 0 {
		query = query.Where("wl.api_type IN ?", req.ApiTypes)
	}
	if req.StartTime != "" {
		query = query.Where("wl.created_at > ?", strings.TrimSpace(req.StartTime))
	}
	if req.EndTime != "" {
		query = query.Where("wl.created_at < ?", strings.TrimSpace(req.EndTime))
	}

	return query
}


func (r *wafLogRepository) GetWafLog(ctx context.Context, id int64) (*model.WafLog, error) {
	var res model.WafLog
	return &res, r.DBWithName(ctx,"admin").Where("id = ?", id).First(&res).Error
}

func (r *wafLogRepository) GetWafLogList(ctx context.Context, req adminApi.SearchWafLogParams) (*v1.PaginatedResponse[model.WafLog], error) {
	var res []model.WafLog
	var total int64

	query := r.DBWithName(ctx,"admin").Model(&model.WafLog{})
	if  req.RequestIp != "" {
		trimmedName := strings.TrimSpace(req.RequestIp)
		query = query.Where("request_ip LIKE CONCAT('%', ?, '%')", trimmedName)
	}


	if req.Uid != 0 {
		query = query.Where("uid = ?", req.Uid)
	}

	if req.Api != "" {
		trimmedName := strings.TrimSpace(req.Api)
		query = query.Where("api LIKE CONCAT('%', ?, '%')", trimmedName)
	}

	if req.Name != "" {
		trimmedName := strings.TrimSpace(req.Name)
		query = query.Where("name LIKE CONCAT('%', ?, '%')", trimmedName)
	}


	if req.RuleId != 0 {
		query = query.Where("rule_id = ?", req.RuleId)
	}

	if req.HostId != 0 {
		query = query.Where("host_id = ?", req.HostId)
	}

	if req.Api != "" {
		trimmedName := strings.TrimSpace(req.Api)
		query = query.Where("api LIKE CONCAT('%', ?, '%')", trimmedName)
	}

	if req.UserAgent != "" {
		trimmedName := strings.TrimSpace(req.UserAgent)
		query = query.Where("user_agent LIKE CONCAT('%', ?, '%')", trimmedName)
	}

	if req.ApiName != "" {
		trimmedName := strings.TrimSpace(req.ApiName)
		query = query.Where("api_name LIKE CONCAT('%', ?, '%')", trimmedName)
	}

	if req.ApiType != "" {
		query = query.Where("api_type = ?", req.ApiType)
	}


	if req.Column != "" {
		query = query.Order(req.Column + " " + req.Order)
	}

	if err := query.Count(&total).Error; err != nil {
		return nil, err
	}


	page := req.Current
	pageSize := req.PageSize

	if page <= 0 {
		page = 1
	}

	if pageSize <= 0 {
		pageSize = 10 
	} else if pageSize > 100 {
		pageSize = 100 
	}

	offset := (page - 1) * pageSize
	result := query.Offset(offset).Limit(pageSize).Find(&res)
	if result.Error != nil {
		return nil, result.Error
	}
	return &v1.PaginatedResponse[model.WafLog]{
		Records: res,
		Page: page,
		PageSize: pageSize,
		Total: total,
		TotalPages: int(math.Ceil(float64(total) / float64(pageSize))),

	}, nil
}

func (r *wafLogRepository) AddWafLog(ctx context.Context, log *model.WafLog) error {
	return r.DBWithName(ctx,"admin").Create(log).Error
}

func (r *wafLogRepository) BatchAddWafLog(ctx context.Context, logs []*model.WafLog) error {
	if len(logs) == 0 {
		return nil
	}
	return r.DBWithName(ctx, "admin").CreateInBatches(logs, len(logs)).Error
}

func (r *wafLogRepository) ExportWafLog(ctx context.Context, req adminApi.ExportWafLog) ([]model.WafLogWithGatewayIP, error) {
	return r.ExportWafLogWithPagination(ctx, req, 0, 0)
}

// ExportWafLogWithPagination 使用子查询获取每条日志在当时时间点的正确网关组IP
func (r *wafLogRepository) ExportWafLogWithPagination(ctx context.Context, req adminApi.ExportWafLog, page, pageSize int) ([]model.WafLogWithGatewayIP, error) {
	var res []model.WafLogWithGatewayIP
	
	// 1. 使用辅助函数构建基础查询
	query := r.buildExportQuery(ctx, req)

	// 2. 构建子查询
	subQuery := r.DBWithName(ctx, "admin").Model(&model.WafLog{}).
		Select("extra_data").
		Where("api_name = ?", "分配网关组").
		Where("host_id = wl.host_id").
		Where("uid = wl.uid").
		Where("created_at <= wl.created_at").
		Order("created_at DESC").
		Limit(1)

	// 3. 添加 Select 和分页
	query = query.Select("wl.*, (?) as gateway_ip_data", subQuery)
	if page > 0 && pageSize > 0 {
		offset := (page - 1) * pageSize
		query = query.Offset(offset).Limit(pageSize)
	}

	// 4. 执行查询
	if err := query.Find(&res).Error; err != nil {
		if err == gorm.ErrRecordNotFound {
			return []model.WafLogWithGatewayIP{}, nil
		}
		return nil, err
	}

	return res, nil
}


// GetWafLogExportCount 获取导出数据总数（已优化）
func (r *wafLogRepository) GetWafLogExportCount(ctx context.Context, req adminApi.ExportWafLog) (int, error) {
	var count int64
	
	// 直接复用 buildExportQuery 来构建查询
	query := r.buildExportQuery(ctx, req)
	
	if err := query.Count(&count).Error; err != nil {
		return 0, err
	}
	
	return int(count), nil
}