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/get":      "获取web详情",
	"/webForward/getList":  "获取web列表",
	"/webForward/add":      "添加web",
	"/webForward/edit":     "修改web",
	"/webForward/delete":   "删除web",

	"/tcpForward/add":      "添加tcp",
	"/tcpForward/edit":     "修改tcp",
	"/tcpForward/delete":   "删除tcp",
	"/tcpForward/getList":  "获取tcp列表",
	"/tcpForward/get":      "获取tcp详情",

	"/udpForward/add":      "添加udp",
	"/udpForward/edit":     "修改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白名单",

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

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,
) 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
}
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,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,
		})
	}
	
	// 调用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))
	}
}

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
	}

	// 使用优化后的转换方法，避免N+1查询
	return s.convertRawDataToExportResults(ctx, data)
}

// SmartExportWafLog 智能导出WAF日志为Excel
func (s *wafLogService) SmartExportWafLog(ctx context.Context, req adminApi.ExportWafLog, w http.ResponseWriter) error {
	// 1. 先获取总数量用于智能选择传输方式
	count, err := s.wafLogRepository.GetWafLogExportCount(ctx, req)
	if err != nil {
		return fmt.Errorf("获取导出数据总数失败: %w", err)
	}

	// 2. 智能选择导出方式
	// 估算每行数据大小约200字节（包含用户名、IP、API名称、域名等字段）
	exportType := excel.SmartExport(count, 200)
	
	// 3. 设置Excel表头映射
	headers := []string{"name", "request_ip", "host_id", "api_name", "addr_backend_list", "domain", "comment", "custom_host", "expose_addr", "created_at"}
	headerMap := map[string]string{
		"name":             "用户名",
		"request_ip":       "请求IP",
		"host_id":          "主机ID", 
		"api_name":         "API名称",
		"addr_backend_list": "后端地址",
		"domain":           "域名",
		"comment":          "备注",
		"custom_host":      "回源地址",
		"expose_addr":      "暴露地址",
		"created_at":       "创建时间",
	}

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

	// 5. 根据导出类型选择不同的处理方式
	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)
	}
}

// normalExportWafLog 普通导出（小文件）
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,
			"api_name":         item.ApiName,
			"addr_backend_list": s.formatBackendList(item.AddrBackendList),
			"domain":           item.Domain,
			"comment":          item.Comment,
			"custom_host":      item.CustomHost,
			"expose_addr":      s.formatExposeAddr(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",
	})
}

// streamExportWafLog 流式导出（大文件）
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")

	// 分批处理数据，每批1000条
	pageSize := 1000
	page := 1

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

		// 转换为导出格式（复用原有的ExPortWafLog逻辑）
		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,
				"api_name":         item.ApiName,
				"addr_backend_list": s.formatBackendList(item.AddrBackendList),
				"domain":           item.Domain,
				"comment":          item.Comment,
				"custom_host":      item.CustomHost,
				"expose_addr":      s.formatExposeAddr(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",
	})
}

// chunkExportWafLog 分块导出（超大文件）
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 // 每个分块5000条记录
	
	// 分块导出需要前端配合实现
	excel.ChunkExport(w, excel.TransferOption{
		FileName:    fileName,
		ContentType: "application/json", // 返回分块信息
	}, totalRecords, pageSize)
	
	return nil
}

// formatBackendList 格式化后端地址列表
func (s *wafLogService) formatBackendList(backendList interface{}) string {
	if backendList == nil {
		return ""
	}
	
	switch v := backendList.(type) {
	case string:
		return v
	case []string:
		return strings.Join(v, ", ")
	default:
		return fmt.Sprintf("%v", v)
	}
}

// formatExposeAddr 格式化暴露地址
func (s *wafLogService) formatExposeAddr(exposeAddr []string) string {
	if len(exposeAddr) == 0 {
		return ""
	}
	return strings.Join(exposeAddr, ", ")
}



// convertRawDataToExportResults 将原始数据转换为导出结果（复用原有的ExPortWafLog逻辑）
func (s *wafLogService) convertRawDataToExportResults(ctx context.Context, rawData []model.WafLog) ([]adminApi.ExportWafLogRes, error) {
	if len(rawData) == 0 {
		return []adminApi.ExportWafLogRes{}, nil
	}

	// 收集所有需要查询的hostId和uid，用于批量获取网关组
	hostIds := make([]int64, 0, len(rawData))
	uids := make([]int64, 0, len(rawData))
	maxCreatedAt := time.Time{}
	
	for _, v := range rawData {
		hostIds = append(hostIds, int64(v.HostId))
		uids = append(uids, int64(v.Uid))
		if v.CreatedAt.After(maxCreatedAt) {
			maxCreatedAt = v.CreatedAt
		}
	}

	// 批量获取网关组数据
	gatewayMap, err := s.wafLogRepository.BatchGetWafLogGateWayIps(ctx, hostIds, uids, maxCreatedAt)
	if err != nil {
		s.Logger.Warn("批量获取网关组失败，降级为单个查询", zap.Error(err))
		gatewayMap = make(map[string]model.WafLog) // 空map，后续会降级处理
	}

	var res []adminApi.ExportWafLogRes
	for _, v := range rawData {
		var AddrBackendList interface{}
		var customHost string
		var port string
		var domain string
		var comment string

		var mapData map[string]interface{}
		err := json.Unmarshal(v.ExtraData, &mapData)
		if err != nil {
			// 尝试解析为数组格式
			var arrayData []interface{}
			if arrayErr := json.Unmarshal(v.ExtraData, &arrayData); arrayErr != nil {
				// 如果不是符合的JSON格式，直接把原始值作为字符串处理
				s.Logger.Warn("额外数据不是有效JSON格式，使用原始值", zap.Error(err), zap.Int("id", v.Id), 
					zap.String("extra_data", string(v.ExtraData)))
				mapData = map[string]interface{}{
					"raw_data": string(v.ExtraData),
				}
			} else {
				// 如果是数组格式，将数组作为值存储
				s.Logger.Warn("额外数据为数组格式，保存为数组值", zap.Int("id", v.Id))
				mapData = map[string]interface{}{
					"array_data": arrayData,
				}
			}
		}

		if strings.Contains(v.ApiName, "tcp") || strings.Contains(v.ApiName, "udp") || strings.Contains(v.ApiName, "web") {
			// 安全地获取extraData
			var extraData map[string]interface{}
			if mapData["data"] != nil {
				if data, ok := mapData["data"].(map[string]interface{}); ok {
					extraData = data
				}
			}
			
			if extraData != nil {
				if extraData["port"] != nil {
					if portStr, ok := extraData["port"].(string); ok {
						port = portStr
					}
				}
				if extraData["domain"] != nil {
					if domainStr, ok := extraData["domain"].(string); ok {
						domain = domainStr
					}
				}
				if extraData["backendList"] != nil {
					if strings.Contains(v.ApiName, "web") {
						if backendListStr, ok := extraData["backendList"].(string); ok {
							var backendList []map[string]interface{}
							err := json.Unmarshal([]byte(backendListStr), &backendList)
							if err != nil {
								s.Logger.Error("解析后端列表失败", zap.Error(err))
								continue
							}
							for _, backend := range backendList {
								if backend["addr"] != nil {
									AddrBackendList = backend["addr"]
								}
								if backend["customHost"] != nil {
									if customHostStr, ok := backend["customHost"].(string); ok {
										customHost = customHostStr
									}
								}
							}
						}
					} else {
						AddrBackendList = extraData["backendList"]
					}
				}
			}
		}

		if mapData["comment"] != nil {
			if commentStr, ok := mapData["comment"].(string); ok {
				comment = commentStr
			}
		}

		// 优化：从批量获取的网关组数据中查找
		var exposeAddr []string
		key := fmt.Sprintf("%d_%d", v.HostId, v.Uid)
		if gatewayModel, exists := gatewayMap[key]; exists {
			var gateWayIps []string
			if err := json.Unmarshal(gatewayModel.ExtraData, &gateWayIps); err == nil {
				if len(gateWayIps) > 0 && port != "" {
					for _, ip := range gateWayIps {
						exposeAddr = append(exposeAddr, ip+":"+port)
					}
				}
			}
		} else {
			// 降级：单个查询
			gateWayIpModel, err := s.wafLogRepository.GetWafLogGateWayIp(ctx, int64(v.HostId), int64(v.Uid), v.CreatedAt)
			if err == nil {
				var gateWayIps []string
				if err := json.Unmarshal(gateWayIpModel.ExtraData, &gateWayIps); err == nil {
					if len(gateWayIps) > 0 && port != "" {
						for _, ip := range gateWayIps {
							exposeAddr = append(exposeAddr, ip+":"+port)
						}
					}
				}
			}
		}

		res = append(res, adminApi.ExportWafLogRes{
			Name:            v.Name,
			RequestIp:       v.RequestIp,
			HostId:          v.HostId,
			ApiName:         v.ApiName,
			AddrBackendList: AddrBackendList,
			Domain:          domain,
			Comment:         comment,
			CustomHost:      customHost,
			ExposeAddr:      exposeAddr,
			CreatedAt:       v.CreatedAt,
		})
	}

	return res, nil
}

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