go_gameshield/internal/service/wafformatter.go

package service

import (
	"context"
	"encoding/json"
	"fmt"
	v1 "github.com/go-nunu/nunu-layout-advanced/api/v1"
	"github.com/go-nunu/nunu-layout-advanced/internal/model"
	"github.com/go-nunu/nunu-layout-advanced/internal/repository"
	"github.com/go-nunu/nunu-layout-advanced/pkg/rabbitmq"
	amqp "github.com/rabbitmq/amqp091-go"
	"go.uber.org/zap"
	"golang.org/x/net/idna"
	"golang.org/x/net/publicsuffix"
	"golang.org/x/sync/errgroup"
	"net"
	"slices"
	"strconv"
	"strings"
)

type WafFormatterService interface {
	Require(ctx context.Context, req v1.GlobalRequire) (RequireResponse, error)
	validateWafPortCount(ctx context.Context, hostId int) error
	validateWafDomainCount(ctx context.Context, req v1.GlobalRequire) error
	ConvertToWildcardDomain(ctx context.Context, domain string) (string, error)
	AppendWafIp(ctx context.Context, req []string, returnSourceIp string) ([]v1.IpInfo, error)
	WashIps(ctx context.Context, req []string) ([]string, error)
	PublishIpWhitelistTask(ips []string, action string, returnSourceIp string, color string)
	PublishDomainWhitelistTask(domain, ip, action string)
	findIpDifferences(oldIps, newIps []string) ([]string, []string)
	WashDeleteWafIp(ctx context.Context, backendList []string) ([]string, error)
	WashEditWafIp(ctx context.Context, newBackendList []string, oldBackendList []string) ([]string, []string, error)
	//cdn添加网站
	AddOrigin(ctx context.Context, req v1.WebJson) (int64, error)
	// 获取ip数量等于1的源站过白ip
	WashDelIps(ctx context.Context, ips []string) ([]string, error)
	// 判断域名是否是IDN,如果是，转换为 Punycode
	ConvertToPunycodeIfIDN(ctx context.Context, domain string) (isIDN bool, punycodeDomain string, err error)
	// 验证端口重复
	VerifyPort(ctx context.Context,protocol string, id int64, port string,hostId int64,domain string) error
}

func NewWafFormatterService(
	service *Service,
	globalRep repository.GlobalLimitRepository,
	hostRep repository.HostRepository,
	required RequiredService,
	parser ParserService,
	tcpforwardingRep repository.TcpforwardingRepository,
	udpForWardingRep repository.UdpForWardingRepository,
	webForwardingRep repository.WebForwardingRepository,
	mq *rabbitmq.RabbitMQ,
	host HostService,
	gatewayGroupRep repository.GatewayGroupRepository,
	gatewayGroupIpRep repository.GateWayGroupIpRepository,
	cdn CdnService,
) WafFormatterService {
	return &wafFormatterService{
		Service:           service,
		globalRep:         globalRep,
		hostRep:           hostRep,
		required:          required,
		parser:            parser,
		tcpforwardingRep:  tcpforwardingRep,
		udpForWardingRep:  udpForWardingRep,
		webForwardingRep:  webForwardingRep,
		host:              host,
		mq:                mq,
		gatewayGroupRep:   gatewayGroupRep,
		gatewayGroupIpRep: gatewayGroupIpRep,
		cdn:               cdn,
	}
}

type wafFormatterService struct {
	*Service
	globalRep         repository.GlobalLimitRepository
	hostRep           repository.HostRepository
	required          RequiredService
	parser            ParserService
	tcpforwardingRep  repository.TcpforwardingRepository
	udpForWardingRep  repository.UdpForWardingRepository
	webForwardingRep  repository.WebForwardingRepository
	host              HostService
	mq                *rabbitmq.RabbitMQ
	gatewayGroupRep   repository.GatewayGroupRepository
	gatewayGroupIpRep repository.GateWayGroupIpRepository
	cdn               CdnService
}

type RequireResponse struct {
	model.GlobalLimit `json:"globalLimit" form:"globalLimit"`
	GatewayIps        []string `json:"ips" form:"ips"`
	Tag               string   `json:"tag" form:"tag"`
}

func (s *wafFormatterService) Require(ctx context.Context, req v1.GlobalRequire) (RequireResponse, error) {
	var res RequireResponse
	// 获取全局配置信息
	globalLimit, err := s.globalRep.GetGlobalLimitByHostId(ctx, int64(req.HostId))
	if err != nil {
		return RequireResponse{}, err
	}
	if globalLimit != nil {
		res.GlobalLimit = *globalLimit
	}
	// 获取主机名
	domain, err := s.hostRep.GetDomainById(ctx, req.HostId)
	if err != nil {
		return RequireResponse{}, err
	}
	res.Tag = strconv.Itoa(req.Uid) + "_" + strconv.Itoa(req.HostId) + "_" + domain + "_" + req.Comment
	res.GatewayIps, err = s.gatewayGroupIpRep.GetGateWayGroupAllIpByGatewayGroupId(ctx, res.GatewayGroupId)
	if err != nil {
		return RequireResponse{}, err
	}
	return res, nil
}

func (s *wafFormatterService) validateWafPortCount(ctx context.Context, hostId int) error {
	congfig, err := s.host.GetGlobalLimitConfig(ctx, hostId)
	if err != nil {
		return err
	}
	tcpCount, err := s.tcpforwardingRep.GetTcpForwardingPortCountByHostId(ctx, hostId)
	if err != nil {
		return err
	}
	udpCount, err := s.udpForWardingRep.GetUdpForwardingPortCountByHostId(ctx, hostId)
	if err != nil {
		return err
	}
	webCount, err := s.webForwardingRep.GetWebForwardingPortCountByHostId(ctx, hostId)
	if err != nil {
		return err
	}
	if int64(congfig.PortCount) > tcpCount+udpCount+webCount {
		return nil
	}
	return fmt.Errorf("端口数量超出套餐限制,已配置%d个端口，套餐限制为%d个端口", tcpCount+udpCount+webCount, congfig.PortCount)
}

func (s *wafFormatterService) validateWafDomainCount(ctx context.Context, req v1.GlobalRequire) error {
	congfig, err := s.host.GetGlobalLimitConfig(ctx, req.HostId)
	if err != nil {
		return err
	}
	domainCount, domainSlice, err := s.webForwardingRep.GetWebForwardingDomainCountByHostId(ctx, req.HostId)
	if err != nil {
		return err
	}
	if req.Domain != "" {
		if !slices.Contains(domainSlice, req.Domain) {
			domainCount += 1
			if domainCount > int64(congfig.DomainCount) {
				return fmt.Errorf("域名数量已达到上限,已配置%d个域名，套餐限制为%d个域名", domainCount, congfig.DomainCount)
			}
		}
	}
	return nil
}
func (s *wafFormatterService) ConvertToWildcardDomain(ctx context.Context, domain string) (string, error) {
	// 1. 使用 EffectiveTLDPlusOne 获取可注册域名部分。
	//    例如，对于 "www.google.com"，这将返回 "google.com"。
	//    对于 "a.b.c.tokyo.jp"，这将返回 "c.tokyo.jp"。
	if domain == "" {
		return "", nil
	}
	registrableDomain, err := publicsuffix.EffectiveTLDPlusOne(domain)
	if err != nil {
		s.logger.Error("无效的域名", zap.String("domain", domain), zap.Error(err))
		// 如果域名无效（如 IP 地址、localhost），则返回错误。
		return "", nil
	}

	// 2. 比较原始域名和可注册域名。
	//    如果它们不相等，说明原始域名包含子域名。
	if domain != registrableDomain {
		// 3. 如果存在子域名，则用 "*." 加上可注册域名来构造通配符域名。
		return registrableDomain, nil
	}

	// 4. 如果原始域名和可注册域名相同（例如，输入就是 "google.com"），
	//    则说明没有子域名可替换，直接返回原始域名。
	return domain, nil
}

func (s *wafFormatterService) AppendWafIp(ctx context.Context, req []string, returnSourceIp string) ([]v1.IpInfo, error) {
	var ips []v1.IpInfo
	for _, v := range req {
		ips = append(ips, v1.IpInfo{
			FType:     "0",
			FStartIp:  v,
			FEndIp:    v,
			FRemark:   "宁波高防IP过白",
			FServerIp: returnSourceIp,
		})
	}
	return ips, nil
}

func (s *wafFormatterService) AppendWafIpByRemovePort(ctx context.Context, req []string) ([]v1.IpInfo, error) {
	var ips []v1.IpInfo
	for _, v := range req {
		ip, _, err := net.SplitHostPort(v)
		if err != nil {
			return nil, err
		}
		ips = append(ips, v1.IpInfo{
			FType:     "0",
			FStartIp:  ip,
			FEndIp:    ip,
			FRemark:   "宁波高防IP过白",
			FServerIp: "",
		})
	}
	return ips, nil

}

func (s *wafFormatterService) WashIps(ctx context.Context, req []string) ([]string, error) {
	var res []string
	for _, v := range req {
		res = append(res, v)
	}
	return res, nil
}

// publishDomainWhitelistTask is a helper function to publish domain whitelist tasks to RabbitMQ.
// It can handle different actions like "add" or "del".
func (s *wafFormatterService) PublishDomainWhitelistTask(domain, ip, action string) {
	// Define message payload, including the action
	type domainTaskPayload struct {
		Domain string `json:"domain"`
		Ip     string `json:"ip"`
		Action string `json:"action"`
	}
	payload := domainTaskPayload{
		Domain: domain,
		Ip:     ip,
		Action: action,
	}

	// Serialize the message
	msgBody, err := json.Marshal(payload)
	if err != nil {
		s.logger.Error("Failed to serialize domain whitelist task message", zap.Error(err), zap.String("domain", domain), zap.String("ip", ip), zap.String("action", action))
		return
	}

	// Get task configuration
	taskCfg, ok := s.mq.GetTaskConfig("domain_whitelist")
	if !ok {
		s.logger.Error("Failed to get 'domain_whitelist' task configuration")
		return
	}

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

	// 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("发布 域名 白名单任务到 MQ 失败", zap.Error(err), zap.String("domain", domain), zap.String("action", action))
	} else {
		s.logger.Info("成功将 域名 白名单任务发布到 MQ", zap.String("domain", domain), zap.String("action", action))
	}
}

func (s *wafFormatterService) PublishIpWhitelistTask(ips []string, action string, returnSourceIp string, color string) {
	// Define message payload, including the action
	type ipTaskPayload struct {
		Ips            []string `json:"ips"`
		Action         string   `json:"action"`
		ReturnSourceIp string   `json:"return_source_ip"`
		Color          string   `json:"color"`
	}
	payload := ipTaskPayload{
		Ips:            ips,
		Action:         action,
		ReturnSourceIp: returnSourceIp,
		Color:          color,
	}

	// Serialize the message
	msgBody, err := json.Marshal(payload)
	if err != nil {
		s.logger.Error("序列化 IP 白名单任务消息失败", zap.Error(err), zap.Any("IPs", ips), zap.String("action", action), zap.String("color", color))
		return
	}

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

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

	// 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("发布 IP 白名单任务到 MQ 失败", zap.Error(err), zap.String("action", action), zap.String("color", color))
	} else {
		s.logger.Info("成功将 IP 白名单任务发布到 MQ", zap.String("action", action), zap.String("color", color))
	}
}

func (s *wafFormatterService) findIpDifferences(oldIps, newIps []string) ([]string, []string) {
	// 使用 map 实现 set，用于快速查找
	oldIpsSet := make(map[string]struct{}, len(oldIps))
	for _, ip := range oldIps {
		oldIpsSet[ip] = struct{}{}
	}

	newIpsSet := make(map[string]struct{}, len(newIps))
	for _, ip := range newIps {
		newIpsSet[ip] = struct{}{}
	}

	var addedIps []string
	// 查找新增的 IP：存在于 newIpsSet 但不存在于 oldIpsSet
	for ip := range newIpsSet {
		if _, found := oldIpsSet[ip]; !found {
			addedIps = append(addedIps, ip)
		}
	}

	var removedIps []string
	// 查找移除的 IP：存在于 oldIpsSet 但不存在于 newIpsSet
	for ip := range oldIpsSet {
		if _, found := newIpsSet[ip]; !found {
			removedIps = append(removedIps, ip)
		}
	}

	return addedIps, removedIps
}

func (s *wafFormatterService) WashDeleteWafIp(ctx context.Context, backendList []string) ([]string, error) {
	var res []string
	for _, v := range backendList {
		ip, _, err := net.SplitHostPort(v)
		if err != nil {
			return nil, err
		}
		res = append(res, ip)
	}
	return res, nil
}

func (s *wafFormatterService) WashEditWafIp(ctx context.Context, newBackendList []string, oldBackendList []string) ([]string, []string, error) {
	var oldIps []string
	var newIps []string
	for _, v := range oldBackendList {
		ip, _, err := net.SplitHostPort(v)
		if err != nil {
			return nil, nil, err
		}
		oldIps = append(oldIps, ip)
	}
	if newBackendList != nil {
		for _, v := range newBackendList {
			ip, _, err := net.SplitHostPort(v)
			if err != nil {
				return nil, nil, err
			}
			newIps = append(newIps, ip)
		}
	}
	addedIps, removedIps := s.findIpDifferences(oldIps, newIps)

	return addedIps, removedIps, nil
}

func (s *wafFormatterService) AddOrigin(ctx context.Context, req v1.WebJson) (int64, error) {
	ip, port, err := net.SplitHostPort(req.BackendList)
	if err != nil {
		return 0, fmt.Errorf("无效的后端地址: %s", err)
	}
	addr := v1.Addr{
		Protocol: req.ApiType,
		Host:     ip,
		Port:     port,
	}
	id, err := s.cdn.CreateOrigin(ctx, v1.Origin{
		Addr:                  addr,
		Weight:                10,
		Description:           req.Comment,
		Host:                  req.Host,
		IsOn:                  true,
		TlsSecurityVerifyMode: "auto",
	})
	if err != nil {
		return 0, err
	}
	return id, nil
}

// 获取ip数量等于1的源站过白ip
func (s *wafFormatterService) WashDelIps(ctx context.Context, ips []string) ([]string, error) {
	var udpIpCounts, tcpIpCounts, webIpCounts []v1.IpCountResult
	g, gCtx := errgroup.WithContext(ctx)
	// 1. 查询 IP 的数量
	g.Go(func() error {
		var err error
		udpIpCounts, err = s.udpForWardingRep.GetIpCountByIp(gCtx, ips)
		if err != nil {
			return fmt.Errorf("in udp repository: %w", err)
		}
		return nil
	})

	g.Go(func() error {
		var err error
		tcpIpCounts, err = s.tcpforwardingRep.GetIpCountByIp(gCtx, ips)
		if err != nil {
			return fmt.Errorf("in tcp repository: %w", err)
		}
		return nil
	})

	g.Go(func() error {
		var err error
		webIpCounts, err = s.webForwardingRep.GetIpCountByIp(gCtx, ips)
		if err != nil {
			return fmt.Errorf("in web repository: %w", err)
		}
		return nil
	})

	if err := g.Wait(); err != nil {
		return nil, err
	}

	// 2. 汇总所有计数结果
	totalCountMap := make(map[string]int)
	// 将多个 for 循环合并到一个函数中，可以显得更整洁（可选）
	accumulateCounts := func(counts []v1.IpCountResult) {
		for _, result := range counts {
			totalCountMap[result.Ip] += result.Count
		}
	}
	accumulateCounts(udpIpCounts)
	accumulateCounts(tcpIpCounts)
	accumulateCounts(webIpCounts)

	// 3. 筛选出总引用数小于 2 的 IP
	var ipsToDelist []string
	for _, ip := range ips {
		if totalCountMap[ip] < 2 {
			ipsToDelist = append(ipsToDelist, ip)
		}
	}

	return ipsToDelist, nil
}

// 判断域名是否为 中文域名,如果是，转换为 Punycode
func (s *wafFormatterService) ConvertToPunycodeIfIDN(ctx context.Context, domain string) (isIDN bool, punycodeDomain string, err error) {
	// 使用 idna.ToASCII 将域名转换为 Punycode。
	// 这个函数同时会根据 IDNA 规范验证域名的合法性。
	punycodeDomain, err = idna.ToASCII(domain)
	if err != nil {
		// 如果转换出错，说明域名格式不符合 IDNA 标准。
		return false, "", fmt.Errorf("域名 '%s' 格式无效: %v", domain, err)
	}

	// 判断是否为 IDN 的关键：
	// 比较转换后的 Punycode 域名和原始域名（忽略大小写）。
	// 如果不相等，说明原始域名包含非 ASCII 字符，即为 IDN。
	isIDN = !strings.EqualFold(domain, punycodeDomain)

	return isIDN, punycodeDomain, nil
}


// 验证端口重复
func (s *wafFormatterService) VerifyPort(ctx context.Context,protocol string, id int64, port string,hostId int64,domain string) error {
	errPortInUse := fmt.Errorf("端口 %s 已经被使用，无法添加", port)
	switch protocol {
		case "http", "https":
			domains, err := s.webForwardingRep.GetDomainByHostIdPort(ctx, hostId, port)
			if err != nil {
				return err
			}
			tcpCount, err := s.tcpforwardingRep.GetPortCount(ctx, hostId, port)
			if err != nil {
				return err
			}

			if tcpCount > 0 {
				return errPortInUse
			}

			for _, v := range domains {
				// 防住空域名修改为非空域名报错
				if v.Domain == "" && int64(v.Id) != id {
					return errPortInUse
				}
				if net.ParseIP(v.Domain) != nil {
					return errPortInUse
				}
			}

			// 确保添加新规则时,没有已有域名的规则
			if net.ParseIP(domain) != nil || domain == "" {
				if len(domains) > 0 {
					return errPortInUse
				}
			}

			return nil


		case "tcp":
			count, err := s.tcpforwardingRep.GetPortCount(ctx, hostId, port)
			if err != nil {
				return err
			}
			webCount, err := s.webForwardingRep.GetDomainByHostIdPort(ctx, hostId, port)
			if err != nil {
				return err
			}
			if count + int64(len(webCount)) > 0 {
				return errPortInUse
			}
			return nil
		case "udp":
			count, err := s.udpForWardingRep.GetPortCount(ctx, hostId, port)
			if err != nil {
				return err
			}
			if count > 0 {
				return errPortInUse
			}
			return nil
		default:
			return fmt.Errorf("不支持的协议类型：%s", protocol)
	}
}