go_gameshield/internal/service/admin/waflog.go

package admin

import (
	"context"
	"encoding/json"
	"fmt"
	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"
	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/excel"
	"github.com/go-nunu/nunu-layout-advanced/pkg/rabbitmq"
	amqp "github.com/rabbitmq/amqp091-go"
	"go.uber.org/zap"
	"net/http"
	"strings"
	"time"
)

// ApiDescriptionMap API描述映射
var ApiDescriptionMap = map[string]string{

	"/webForward/add":      "添加web",
	"/webForward/edit":     "修改web",
	"/webForward/delete":   "删除web",

	"/tcpForward/add":      "添加tcp",
	"/tcpForward/edit":     "修改tcp",
	"/tcpForward/delete":   "删除tcp",


	"/udpForward/add":      "添加udp",
	"/udpForward/edit":     "修改udp",
	"/udpForward/delete":   "删除udp",


	"/globalLimit/add":     "添加实例",
	"/globalLimit/edit":    "修改实例",
	"/globalLimit/delete":  "删除实例",


	"/allowAndDeny/add":      "添加黑白名单",
	"/allowAndDeny/edit":     "修改黑白名单",
	"/allowAndDeny/delete":   "删除黑白名单",


	"/cc/editState":  "删除CC黑名单",


	"/ccIpList/add":      "添加CC白名单",
	"/ccIpList/edit":     "修改CC白名单",
	"/ccIpList/delete":   "删除CC白名单",

	"分配网关组": "分配网关组",
}

type WafLogService 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, req adminApi.WafLog) error
	BatchAddWafLog(ctx context.Context, reqs []*adminApi.WafLog) error
	PublishIpWafLogTask(ctx context.Context, req adminApi.WafLog)
	SmartExportWafLog(ctx context.Context, req adminApi.ExportWafLog, w http.ResponseWriter) error
	GetApiDescriptions(ctx context.Context) map[string]string
}
func NewWafLogService(
    service *service.Service,
    wafLogRepository adminRep.WafLogRepository,
	globalLimitRepository waf.GlobalLimitRepository,
	mq *rabbitmq.RabbitMQ,
	wafLogDataCleanService WafLogDataCleanService,
) WafLogService {
	return &wafLogService{
		Service:        service,
		wafLogRepository: wafLogRepository,
		globalLimitRepository: globalLimitRepository,
		mq : mq,
		wafLogDataCleanService: wafLogDataCleanService,
	}
}

type wafLogService struct {
	*service.Service
	wafLogRepository adminRep.WafLogRepository
	globalLimitRepository waf.GlobalLimitRepository
	mq *rabbitmq.RabbitMQ
	wafLogDataCleanService WafLogDataCleanService
}
func (s *wafLogService) getFirstPathSegment(path string) (segment []string, ok bool) {
	trimmedPath := strings.TrimPrefix(path, "/")
	if trimmedPath == "" {
		return nil, false
	}
	parts := strings.Split(trimmedPath, "/")
	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,req adminApi.SearchWafLogParams) (*v1.PaginatedResponse[model.WafLog], error) {
	return s.wafLogRepository.GetWafLogList(ctx, req)
}

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
	}

	req.RequestIp = userInfo.LastLoginIp

	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
		req.RequestIp = userInfo.LastLoginIp
		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,
		})
	}
	
	return s.wafLogRepository.BatchAddWafLog(ctx, wafLogs)
}

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

	payload := &req

	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
	}

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

	routingKey := fmt.Sprintf("wafLog.%s", "add")

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

	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))
	}
}

func (s *wafLogService) ExPortWafLog(ctx context.Context, req adminApi.ExportWafLog) ([]adminApi.ExportWafLogRes, error) {
	data, err := s.wafLogRepository.ExportWafLog(ctx, req)
	if err != nil {
		return nil, err
	}
	return s.convertRawDataToExportResults(ctx, data)
}

func (s *wafLogService) SmartExportWafLog(ctx context.Context, req adminApi.ExportWafLog, w http.ResponseWriter) error {
	count, err := s.wafLogRepository.GetWafLogExportCount(ctx, req)
	if err != nil {
		return fmt.Errorf("获取导出数据总数失败: %w", err)
	}

	exportType := excel.SmartExport(count, 200)
	
	headers := []string{"name", "request_ip", "host_id", "rule_id", "api_name", "addr_backend_list", "allow_and_deny_ips", "domain", "custom_host", "expose_addr", "comment", "created_at"}
	headerMap := map[string]string{
		"name":             "用户名",
		"request_ip":       "请求IP",
		"host_id":          "实例ID",
		"rule_id":          "规则ID",
		"api_name":         "操作名称",
		"addr_backend_list": "后端地址",
		"allow_and_deny_ips": "黑白名单",
		"domain":           "域名",
		"custom_host":      "回源地址",
		"expose_addr":      "暴露地址",
		"comment":          "备注",
		"created_at":       "操作时间",
	}

	generator := excel.NewExcelGenerator("WAF日志", headers, headerMap)
	if err := generator.WriteHeaders(); err != nil {
		return fmt.Errorf("写入Excel表头失败: %w", err)
	}

	switch exportType {
	case excel.ExportTypeNormal:
		return s.normalExportWafLog(ctx, req, generator, w)
	case excel.ExportTypeStream:
		return s.streamExportWafLog(ctx, req, generator, w)
	case excel.ExportTypeChunk:
		return s.chunkExportWafLog(ctx, req, w, count)
	default:
		return s.normalExportWafLog(ctx, req, generator, w)
	}
}

func (s *wafLogService) normalExportWafLog(ctx context.Context, req adminApi.ExportWafLog, generator *excel.ExcelGenerator, w http.ResponseWriter) error {
	exportData, err := s.ExPortWafLog(ctx, req)
	if err != nil {
		return fmt.Errorf("获取导出数据失败: %w", err)
	}

	data := make([]map[string]interface{}, 0, len(exportData))
	for _, item := range exportData {
		row := map[string]interface{}{
			"name":             item.Name,
			"request_ip":       item.RequestIp,
			"host_id":          item.HostId,
			"rule_id":          s.wafLogDataCleanService.FormatBackendList(item.RuleId),
			"api_name":         item.ApiName,
			"addr_backend_list": s.wafLogDataCleanService.FormatBackendList(item.AddrBackendList),
			"allow_and_deny_ips": item.AllowAndDenyIps,
			"domain":           item.Domain,
			"comment":          item.Comment,
			"custom_host":     s.wafLogDataCleanService.FormatBackendList(item.CustomHost),
			"expose_addr":      s.wafLogDataCleanService.FormatBackendList(item.ExposeAddr),
			"created_at":       item.CreatedAt,
		}
		data = append(data, row)
	}

	if err := generator.WriteRows(data); err != nil {
		return fmt.Errorf("写入Excel数据失败: %w", err)
	}

	fileName := fmt.Sprintf("waf_logs_%s.xlsx", time.Now().Format("20060102_150405"))
	return excel.NormalExport(generator, w, excel.TransferOption{
		FileName:    fileName,
		ContentType: "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
	})
}

func (s *wafLogService) streamExportWafLog(ctx context.Context, req adminApi.ExportWafLog, generator *excel.ExcelGenerator, w http.ResponseWriter) error {
	fileName := fmt.Sprintf("waf_logs_%s.xlsx", time.Now().Format("20060102_150405"))
	
	w.Header().Set("Content-Type", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet")
	w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%s", fileName))
	w.Header().Set("Transfer-Encoding", "chunked")

	pageSize := 1000
	page := 1

	for {
		exportData, err := s.wafLogRepository.ExportWafLogWithPagination(ctx, req, page, pageSize)
		if err != nil {
			return fmt.Errorf("获取第%d页数据失败: %w", page, err)
		}

		exportResults, err := s.convertRawDataToExportResults(ctx, exportData)
		if err != nil {
			return fmt.Errorf("转换导出数据失败: %w", err)
		}

		if len(exportResults) == 0 {
			break 
		}

		for _, item := range exportResults {
			row := map[string]interface{}{
				"name":             item.Name,
				"request_ip":       item.RequestIp,
				"host_id":          item.HostId,
				"rule_id":          s.wafLogDataCleanService.FormatBackendList(item.RuleId),
				"api_name":         item.ApiName,
				"addr_backend_list": s.wafLogDataCleanService.FormatBackendList(item.AddrBackendList),
				"allow_and_deny_ips": item.AllowAndDenyIps,
				"domain":           item.Domain,
				"comment":          item.Comment,
				"custom_host":      s.wafLogDataCleanService.FormatBackendList(item.CustomHost),
				"expose_addr":      s.wafLogDataCleanService.FormatBackendList(item.ExposeAddr),
				"created_at":       item.CreatedAt,
			}
			
			if err := generator.WriteRow(row); err != nil {
				return fmt.Errorf("写入第%d页数据失败: %w", page, err)
			}
		}

		if len(exportResults) < pageSize {
			break
		}

		page++
	}

	return excel.StreamExport(generator, w, excel.TransferOption{
		FileName:    fileName,
		ContentType: "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
	})
}

func (s *wafLogService) chunkExportWafLog(ctx context.Context, req adminApi.ExportWafLog, w http.ResponseWriter, totalRecords int) error {
	fileName := fmt.Sprintf("waf_logs_%s.xlsx", time.Now().Format("20060102_150405"))
	pageSize := 5000 
	
	excel.ChunkExport(w, excel.TransferOption{
		FileName:    fileName,
		ContentType: "application/json", 
	}, totalRecords, pageSize)
	
	return nil
}

// convertRawDataToExportResults 将从数据库获取的带有关联IP的数据转换为最终导出格式
func (s *wafLogService) convertRawDataToExportResults(ctx context.Context, rawData []model.WafLogWithGatewayIP) ([]adminApi.ExportWafLogRes, error) {
	if len(rawData) == 0 {
		return []adminApi.ExportWafLogRes{}, nil
	}

	var res []adminApi.ExportWafLogRes
	for _, v := range rawData {
		// --- 数据清洗 ---
		cleanedData := s.wafLogDataCleanService.ParseWafLogExtraData(v.ExtraData, v.ApiName)

		// --- 网关 IP 处理 ---
		var exposeAddr []string
		// 直接使用查询结果中附带的 gateway_ip_data
		if len(v.GatewayIpData) > 0 && string(v.GatewayIpData) != "null" {
			var gateWayIps []string
			// 解析从数据库子查询得到的JSON数据
			if err := json.Unmarshal(v.GatewayIpData, &gateWayIps); err == nil && len(gateWayIps) > 0 && cleanedData.Port != "" {
				for _, ip := range gateWayIps {
					exposeAddr = append(exposeAddr, ip+":"+cleanedData.Port)
				}
			}
		}

		var ruleIds []int64
		if len(cleanedData.RuleID) > 0 {
			ruleIds = cleanedData.RuleID
		} else {
			ruleIds = []int64{int64(v.RuleId)}
		}

		if cleanedData.IsHttps == 1 {
			for i := range exposeAddr {
				exposeAddr[i] = "https://" + exposeAddr[i]
			}
		}

		// --- 构造结果 ---
		res = append(res, adminApi.ExportWafLogRes{
			Name:            v.Name,
			RequestIp:       v.RequestIp,
			HostId:          v.HostId,
			RuleId:          ruleIds,
			ApiName:         v.ApiName,
			AddrBackendList: cleanedData.AddrBackendList,
			AllowAndDenyIps: cleanedData.AllowAndDenyIps,
			Domain:          cleanedData.Domain,
			Comment:         cleanedData.Comment,
			CustomHost:      cleanedData.CustomHost,
			ExposeAddr:      exposeAddr,
			CreatedAt:       v.CreatedAt,
		})
	}

	return res, nil
}

// GetApiDescriptions 获取API描述映射
func (s *wafLogService) GetApiDescriptions(ctx context.Context) map[string]string {
	return ApiDescriptionMap
}