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/repository" "github.com/go-nunu/nunu-layout-advanced/pkg/rabbitmq" amqp "github.com/rabbitmq/amqp091-go" "github.com/spf13/cast" "go.uber.org/zap" "golang.org/x/net/publicsuffix" "net" "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) 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) PublishDomainWhitelistTask(domain, ip, action string) findIpDifferences(oldIps, newIps []string) ([]string, []string) WashDeleteWafIp(ctx context.Context, backendList []string,allowIpList []string) ([]string, error) WashEditWafIp(ctx context.Context, newBackendList []string,newAllowIpList []string,oldBackendList []string,oldAllowIpList []string) ([]string, []string, []string, []string, error) GetIp(ctx context.Context, gatewayGroupId int) ([]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, mq *rabbitmq.RabbitMQ, host HostService, gatewayGroupRep repository.GatewayGroupRepository, gatewayGroupIpRep repository.GateWayGroupIpRepository, ) 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, } } 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 } 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 } 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) { // Define message payload, including the action type ipTaskPayload struct { Ips []string `json:"ips"` Action string `json:"action"` ReturnSourceIp string `json:"return_source_ip"` } payload := ipTaskPayload{ Ips: ips, Action: action, ReturnSourceIp: returnSourceIp, } // 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)) 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)) } else { s.logger.Info("成功将 IP 白名单任务发布到 MQ", zap.String("action", action)) } } 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,allowIpList []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) } res = append(res, allowIpList...) return res, nil } func (s *wafFormatterService) WashEditWafIp(ctx context.Context, newBackendList []string,newAllowIpList []string,oldBackendList []string,oldAllowIpList []string) ([]string, []string, []string, []string, error) { var oldIps []string var newIps []string var oldAllowIps []string var newAllowIps []string for _, v := range oldBackendList { ip, _, err := net.SplitHostPort(v) if err != nil { return nil, nil, 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, nil, nil, err } newIps = append(newIps, ip) } } addedIps, removedIps := s.findIpDifferences(oldIps, newIps) if oldAllowIpList != nil { oldAllowIps = append(oldAllowIps, oldAllowIpList...) } if newAllowIpList != nil { newAllowIps = append(newAllowIps, newAllowIpList...) } addedAllowIps, removedAllowIps := s.findIpDifferences(oldAllowIps, newAllowIps) return addedIps, removedIps ,addedAllowIps, removedAllowIps, nil } func (s *wafFormatterService) GetIp(ctx context.Context, gatewayGroupId int) ([]string,string, error) { WafGatewayGroupRuleId, err := s.gatewayGroupRep.GetGatewayGroupByRuleId(ctx, int64(gatewayGroupId)) if err != nil { return nil, "", err } ips, err := s.gatewayGroupIpRep.GetGateWayGroupAllIpByGatewayGroupId(ctx, WafGatewayGroupRuleId.Id) if err != nil { return nil, "", err } if len(ips) == 0 { return nil, "", fmt.Errorf("请联系客服分配网关IP") } return ips,ips[0], nil }