go_gameshield/internal/service/admin/waflog.go

package admin

import (
	"context"
	"encoding/json"
	"fmt"
	adminApi "github.com/go-nunu/nunu-layout-advanced/api/v1/admin"
	"github.com/go-nunu/nunu-layout-advanced/internal/model"
	adminRep "github.com/go-nunu/nunu-layout-advanced/internal/repository/admin"
	"github.com/go-nunu/nunu-layout-advanced/internal/repository/api/waf"
	"github.com/go-nunu/nunu-layout-advanced/internal/service"
	"github.com/go-nunu/nunu-layout-advanced/pkg/rabbitmq"
	amqp "github.com/rabbitmq/amqp091-go"
	"go.uber.org/zap"
	"strings"
)

type WafLogService interface {
	GetWafLog(ctx context.Context, id int64) (*model.WafLog, error)
	GetWafLogList(ctx context.Context) ([]model.WafLog, error)
	AddWafLog(ctx context.Context, req adminApi.WafLog) error
	BatchAddWafLog(ctx context.Context, reqs []*adminApi.WafLog) error
	PublishIpWafLogTask(ctx context.Context, req adminApi.WafLog)
}
func NewWafLogService(
    service *service.Service,
    wafLogRepository adminRep.WafLogRepository,
	globalLimitRepository waf.GlobalLimitRepository,
	mq *rabbitmq.RabbitMQ,
) WafLogService {
	return &wafLogService{
		Service:        service,
		wafLogRepository: wafLogRepository,
		globalLimitRepository: globalLimitRepository,
		mq : mq,
	}
}

type wafLogService struct {
	*service.Service
	wafLogRepository adminRep.WafLogRepository
	globalLimitRepository waf.GlobalLimitRepository
	mq *rabbitmq.RabbitMQ
}

var ApiDescriptionMap = map[string]string{

	"/webForward/get": "获取web详情",
	"/webForward/getList" : "获取web列表",
	"/webForward/add" : "添加web",
	"/webForward/update" : "更新web",
	"/webForward/delete" : "删除web",

	"/tcpForward/add" : "添加tcp",
	"/tcpForward/update" : "更新tcp",
	"/tcpForward/delete" : "删除tcp",
	"/tcpForward/getList" : "获取tcp列表",
	"/tcpForward/get" : "获取tcp详情",

	"/udpForward/add" : "添加udp",
	"/udpForward/update" : "更新udp",
	"/udpForward/delete" : "删除udp",
	"/udpForward/getList" : "获取udp列表",
	"/udpForward/get" : "获取udp详情",

	"/globalLimit/add" : "添加实例",
	"/globalLimit/edit" : "编辑实例",
	"/globalLimit/delete" : "删除实例",

	"/allowAndDeny/get" : "获取黑白名单详情",
	"/allowAndDeny/getList" : "获取黑白名单列表",
	"/allowAndDeny/add" : "添加黑白名单",
	"/allowAndDeny/edit" : "编辑黑白名单",
	"/allowAndDeny/delete" : "删除黑白名单",

	"/cc/getList" : "获取CC列表",
	"/cc/editState" : "删除CC黑名单",

	"/ccIpList/getList" : "获取CC白名单列表",
	"/ccIpList/add" : "添加CC白名单",
	"/ccIpList/edit" : "编辑CC白名单",
	"/ccIpList/delete" : "删除CC白名单",

}
func (s *wafLogService) getFirstPathSegment(path string) (segment []string, ok bool) {
	// 1. 为了统一处理，先去掉路径最前面的 "/"
	// 这样 "/v1/admin" 会变成 "v1/admin"，而 "v1/admin" 保持不变
	trimmedPath := strings.TrimPrefix(path, "/")

	// 如果去掉 "/" 后字符串为空（比如原路径是 "/" 或 ""），则无法提取
	if trimmedPath == "" {
		return nil, false
	}

	// 2. 使用 "/" 作为分隔符来切割字符串
	// "v1/admin/menus" 会被切割成一个字符串切片 (slice): ["v1", "admin", "menus"]
	parts := strings.Split(trimmedPath, "/")

	// 3. 只要切片不为空，第一个元素就是我们需要的值
	// len(parts) > 0 这个检查可以保证程序不会因为空切片而出错
	if len(parts) > 0 {
		return parts, true
	}

	return nil, false
}

func (s *wafLogService) GetWafLog(ctx context.Context, id int64) (*model.WafLog, error) {
	return s.wafLogRepository.GetWafLog(ctx, id)
}

func (s *wafLogService) GetWafLogList(ctx context.Context) ([]model.WafLog, error) {
	return s.wafLogRepository.GetWafLogList(ctx)
}

func (s *wafLogService) AddWafLog(ctx context.Context, req adminApi.WafLog) error {
	if req.Api != "" {
		api := strings.TrimPrefix(req.Api, "/v1")

		if _, ok := ApiDescriptionMap[api]; ok {
			req.ApiName = ApiDescriptionMap[api]
		}

		apiType, ok := s.getFirstPathSegment(req.Api)
		if ok {
			req.ApiType = apiType[len(apiType)-1]
		}

	}
	userInfo, err := s.globalLimitRepository.GetUserInfo(ctx, int64(req.Uid))
	if err != nil {
		return err
	}
	req.Name = userInfo.Username
	extraData, err := json.Marshal(req.ExtraData)
	if err != nil {
		return err
	}
	return s.wafLogRepository.AddWafLog(ctx, &model.WafLog{
		Uid:        req.Uid,
		Name:       req.Name,
		RequestIp:  req.RequestIp,
		RuleId:     req.RuleId,
		HostId:     req.HostId,
		UserAgent:  req.UserAgent,
		Api:        req.Api,
		ApiType:    req.ApiType,
		ApiName:    req.ApiName,
		ExtraData:  extraData,

	})
}

func (s *wafLogService) BatchAddWafLog(ctx context.Context, reqs []*adminApi.WafLog) error {
	if len(reqs) == 0 {
		return nil
	}
	
	wafLogs := make([]*model.WafLog, 0, len(reqs))
	
	for _, req := range reqs {
		if req.Api != "" {
			api := strings.TrimPrefix(req.Api, "/v1")
			if _, ok := ApiDescriptionMap[api]; ok {
				req.ApiName = ApiDescriptionMap[api]
			}
			
			apiType, ok := s.getFirstPathSegment(req.Api)
			if ok {
				req.ApiType = apiType[len(apiType)-1]
			}
		}
		
		userInfo, err := s.globalLimitRepository.GetUserInfo(ctx, int64(req.Uid))
		if err != nil {
			s.Logger.Error("获取用户信息失败", zap.Error(err), zap.Int("uid", req.Uid))
			continue
		}
		
		req.Name = userInfo.Username
		extraData, err := json.Marshal(req.ExtraData)
		if err != nil {
			s.Logger.Error("序列化额外数据失败", zap.Error(err))
			continue
		}
		
		wafLogs = append(wafLogs, &model.WafLog{
			Uid:        req.Uid,
			Name:       req.Name,
			RequestIp:  req.RequestIp,
			RuleId:     req.RuleId,
			HostId:     req.HostId,
			UserAgent:  req.UserAgent,
			Api:        req.Api,
			ApiType:    req.ApiType,
			ApiName:    req.ApiName,
			ExtraData:  extraData,
		})
	}
	
	// 调用repository层的批量插入方法
	return s.wafLogRepository.BatchAddWafLog(ctx, wafLogs)
}

func (s *wafLogService) PublishIpWafLogTask(ctx context.Context, req adminApi.WafLog) {

	payload := &req

	// Serialize the message
	msgBody, err := json.Marshal(payload)
	if err != nil {
		s.Logger.Error("序列化 WafLog 任务消息失败", zap.Error(err), zap.Int("hostId", payload.HostId), zap.Int("uid", payload.Uid), zap.Any("req", req))
		return
	}

	// Get task configuration
	taskCfg, ok := s.mq.GetTaskConfig("waf_log")
	if !ok {
		s.Logger.Error("无法获取“waf_Log”任务配置")
		return
	}

	// Construct the routing key dynamically based on the action
	routingKey := fmt.Sprintf("wafLog.%s", "add")

	// Construct the amqp.Publishing message
	publishingMsg := amqp.Publishing{
		ContentType:  "application/json",
		Body:         msgBody,
		DeliveryMode: amqp.Persistent, // Persistent message
	}

	// Publish the message
	err = s.mq.PublishWithCh(taskCfg.Exchange, routingKey, publishingMsg)
	if err != nil {
		s.Logger.Error("发布 WafLog 任务消息失败", zap.Error(err), zap.Int("hostId", payload.HostId), zap.Int("uid", payload.Uid), zap.Any("req", req))
	} else {
		s.Logger.Info("已成功发布 WafLog 任务消息", zap.Int("hostId", payload.HostId), zap.Int("uid", payload.Uid), zap.Any("req", req))
	}
}