backtrace/bk/utils.go

package backtrace

import (
	"encoding/json"
	"fmt"
	"io"
	"strings"
	"time"

	"github.com/imroc/req/v3"
	"github.com/oneclickvirt/backtrace/model"
	. "github.com/oneclickvirt/defaultset"
)

type Result struct {
	i int
	s string
}

// removeDuplicates 切片去重
func removeDuplicates(elements []string) []string {
	if elements == nil {
		return nil
	}
	seen := make(map[string]struct{})
	var result []string
	for _, v := range elements {
		if _, ok := seen[v]; !ok {
			seen[v] = struct{}{}
			result = append(result, v)
		}
	}
	return result
}

// getData 获取目标地址的文本内容
func getData(endpoint string) string {
	client := req.C()
	client.SetTimeout(6 * time.Second)
	client.R().
		SetRetryCount(2).
		SetRetryBackoffInterval(1*time.Second, 5*time.Second).
		SetRetryFixedInterval(2 * time.Second)
	if model.EnableLoger {
		InitLogger()
		defer Logger.Sync()
	}
	for _, baseUrl := range model.CdnList {
		url := baseUrl + endpoint
		resp, err := client.R().Get(url)
		if err == nil {
			defer resp.Body.Close()
			b, err := io.ReadAll(resp.Body)
			if strings.Contains(string(b), "error") {
				continue
			}
			if err == nil {
				if model.EnableLoger {
					Logger.Info(fmt.Sprintf("Received data length: %d", len(b)))
				}
				return string(b)
			}
		}
		if model.EnableLoger {
			Logger.Info(fmt.Sprintf("HTTP request failed: %v", err))
		}
	}
	return ""
}

// parseIcmpTargets 解析ICMP目标数据
func parseIcmpTargets(jsonData string) []model.IcmpTarget {
	var targets []model.IcmpTarget
	err := json.Unmarshal([]byte(jsonData), &targets)
	if err != nil {
		if model.EnableLoger {
			Logger.Info(fmt.Sprintf("解析ICMP目标失败: %s", err.Error()))
		}
		return nil
	}
	return targets
}

// tryAlternativeIPs 从IcmpTargets获取备选IP地址
func tryAlternativeIPs(targetName string, ipVersion string) []string {
	if model.ParsedIcmpTargets == nil || (model.ParsedIcmpTargets != nil && len(model.ParsedIcmpTargets) == 0) {
		return nil
	}
	if model.EnableLoger {
		Logger.Info(fmt.Sprintf("使用备选地址: %s %s", targetName, ipVersion))
	}
	// 从目标名称中提取省份和ISP信息
	var targetProvince, targetISP string
	if strings.Contains(targetName, "北京") {
		targetProvince = "北京"
	} else if strings.Contains(targetName, "上海") {
		targetProvince = "上海"
	} else if strings.Contains(targetName, "广州") {
		targetProvince = "广东"
	} else if strings.Contains(targetName, "成都") {
		targetProvince = "四川"
	}
	if strings.Contains(targetName, "电信") {
		targetISP = "电信"
	} else if strings.Contains(targetName, "联通") {
		targetISP = "联通"
	} else if strings.Contains(targetName, "移动") {
		targetISP = "移动"
	}
	// 如果没有提取到信息，返回空
	if targetProvince == "" || targetISP == "" {
		return nil
	}
	// 查找匹配条件的目标
	var result []string
	for _, target := range model.ParsedIcmpTargets {
		// 检查省份是否匹配（可能带有"省"字或不带）
		provinceMatch := (target.Province == targetProvince) || (target.Province == targetProvince+"省")
		// 检查ISP和IP版本是否匹配
		if provinceMatch && target.ISP == targetISP && target.IPVersion == ipVersion {
			// 解析IP列表
			if target.IPs != "" {
				ips := strings.Split(target.IPs, ",")
				// 最多返回3个IP地址
				count := 0
				for _, ip := range ips {
					if ip != "" {
						result = append(result, strings.TrimSpace(ip))
						count++
						if count >= 3 {
							break
						}
					}
				}
				if len(result) > 0 {
					return result
				}
			}
		}
	}
	return nil
}

// mergeHops 合并多个hops结果
func mergeHops(allHops [][]*Hop) []*Hop {
	if len(allHops) == 0 {
		return nil
	}
	if len(allHops) == 1 {
		return allHops[0]
	}
	// 找到最大长度
	maxLen := 0
	for _, hops := range allHops {
		if len(hops) > maxLen {
			maxLen = len(hops)
		}
	}
	var mergedHops []*Hop
	// 逐位置合并
	for pos := 0; pos < maxLen; pos++ {
		var availableHops []*Hop
		// 收集当前位置所有可用的hop
		for _, hops := range allHops {
			if pos < len(hops) {
				availableHops = append(availableHops, hops[pos])
			}
		}
		if len(availableHops) == 0 {
			continue
		}
		// 如果只有一个可用hop，直接使用
		if len(availableHops) == 1 {
			mergedHops = append(mergedHops, availableHops[0])
			continue
		}
		// 统计相同的hop（通过比较第一个node的IP）
		hopCount := make(map[string][]*Hop)
		for _, hop := range availableHops {
			var key string
			if len(hop.Nodes) > 0 && hop.Nodes[0].IP != nil {
				key = hop.Nodes[0].IP.String()
			} else {
				key = "*" // 超时或无响应
			}
			if _, exists := hopCount[key]; !exists {
				hopCount[key] = make([]*Hop, 0)
			}
			hopCount[key] = append(hopCount[key], hop)
		}
		// 按多数原则选择hop
		if len(hopCount) == 1 {
			// 所有hop都相同，选择第一个
			mergedHops = append(mergedHops, availableHops[0])
		} else {
			// 找出最多的hop类型
			maxCount := 0
			var majorityHops []*Hop
			for _, hops := range hopCount {
				if len(hops) > maxCount {
					maxCount = len(hops)
					majorityHops = hops
				}
			}
			// 如果有多数派，使用多数派的第一个
			if maxCount > 1 || len(hopCount) == 2 {
				mergedHops = append(mergedHops, majorityHops[0])
			} else {
				// 三个都不同，按请求早晚顺序选择第一个
				mergedHops = append(mergedHops, availableHops[0])
			}
		}
	}
	return mergedHops
}