go_gameshield/internal/service/wafformatter.go

package service

import (
	"context"
	"fmt"
	v1 "github.com/go-nunu/nunu-layout-advanced/api/v1"
	"github.com/go-nunu/nunu-layout-advanced/internal/repository"
	"github.com/spf13/cast"
	"golang.org/x/net/publicsuffix"
	"slices"
	"strconv"
)

type WafFormatterService interface {
	require(ctx context.Context, req v1.GlobalRequire, category string) (v1.GlobalRequire, error)
	sendFormData(ctx context.Context,addTokenUrl string,addSendUrl string,formData map[string]interface{}) (int, error)
	validateWafPortCount(ctx context.Context, hostId int) error
	validateWafDomainCount(ctx context.Context, req v1.GlobalRequire) error
	ConvertToWildcardDomain(ctx context.Context,domain string) (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,
	host HostService,
) WafFormatterService {
	return &wafFormatterService{
		Service:        service,
		globalRep: globalRep,
		hostRep: hostRep,
		required: required,
		parser: parser,
		tcpforwardingRep: tcpforwardingRep,
		udpForWardingRep: udpForWardingRep,
		webForwardingRep: webForwardingRep,
		host : host,
	}
}

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
}

func (s *wafFormatterService) require(ctx context.Context,req v1.GlobalRequire,category string) (v1.GlobalRequire, error) {
	RuleIds, err := s.globalRep.GetGlobalLimitByHostId(ctx, int64(req.HostId))
	if err != nil {
		return v1.GlobalRequire{}, err
	}
	req.WafGatewayGroupId = RuleIds.GatewayGroupId
	switch category {
	case "tcp":
		req.LimitRuleId = RuleIds.TcpLimitRuleId
	case "udp":
		req.LimitRuleId = RuleIds.UdpLimitRuleId
	case "web":
		req.LimitRuleId = RuleIds.WebLimitRuleId
	}
	domain, err := s.hostRep.GetDomainById(ctx, req.HostId)
	if err != nil {
		return v1.GlobalRequire{}, err
	}
	req.Tag = strconv.Itoa(req.Uid) + "_" + strconv.Itoa(req.HostId) + "_" + domain + "_" + req.Comment
	return req, nil
}

func (s *wafFormatterService) sendFormData(ctx context.Context,addTokenUrl string,addSendUrl string,formData map[string]interface{}) (int, error) {
	respBody, err := s.required.SendForm(ctx, addTokenUrl, addSendUrl, formData)
	if err != nil {
		return 0, err
	}
	// 解析响应内容中的 alert 消息
	res, err := s.parser.ParseAlert(string(respBody))
	if err != nil {
		return 0,err
	}
	if res != "" {
		return 0,fmt.Errorf(res)
	}
	ruleIdStr, err := s.parser.GetRuleIdByColumnName(ctx, respBody,formData["tag"].(string))
	if err != nil {
		return 0, err
	}
	ruleId, err := cast.ToIntE(ruleIdStr)
	if err != nil {
		return 0,err
	}
	return ruleId, 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"。
	registrableDomain, err := publicsuffix.EffectiveTLDPlusOne(domain)
	if err != nil {
		// 如果域名无效（如 IP 地址、localhost），则返回错误。
		return "", fmt.Errorf("无法处理 '%s': %w", domain, err)
	}

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

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