fscan/Common/config/PortMapping.go

package config

import (
	"sort"
	"sync"
)

// ProbeMapping 探测器映射管理器
type ProbeMapping struct {
	mu          sync.RWMutex
	defaultMap  []string
	portMap     map[int][]string
	initialized bool
}

// NewProbeMapping 创建探测器映射管理器
func NewProbeMapping() *ProbeMapping {
	return &ProbeMapping{
		defaultMap:  getDefaultProbeMap(),
		portMap:     getDefaultPortMap(),
		initialized: true,
	}
}

// getDefaultProbeMap 获取默认的探测器顺序
func getDefaultProbeMap() []string {
	// 返回常量的副本
	result := make([]string, len(DefaultProbeMap))
	copy(result, DefaultProbeMap)
	return result
}

// getDefaultPortMap 获取默认的端口映射
func getDefaultPortMap() map[int][]string {
	// 返回常量的深拷贝
	result := make(map[int][]string)
	for port, probes := range DefaultPortMap {
		probesCopy := make([]string, len(probes))
		copy(probesCopy, probes)
		result[port] = probesCopy
	}
	return result
}

// GetProbesForPort 获取指定端口的探测器列表
func (pm *ProbeMapping) GetProbesForPort(port int) []string {
	pm.mu.RLock()
	defer pm.mu.RUnlock()

	if probes, exists := pm.portMap[port]; exists {
		// 返回副本，避免外部修改
		result := make([]string, len(probes))
		copy(result, probes)
		return result
	}

	// 返回默认探测器
	result := make([]string, len(pm.defaultMap))
	copy(result, pm.defaultMap)
	return result
}

// SetProbesForPort 设置指定端口的探测器列表
func (pm *ProbeMapping) SetProbesForPort(port int, probes []string) {
	pm.mu.Lock()
	defer pm.mu.Unlock()

	if len(probes) > 0 {
		// 创建副本
		probesCopy := make([]string, len(probes))
		copy(probesCopy, probes)
		pm.portMap[port] = probesCopy
	} else {
		delete(pm.portMap, port)
	}
}

// AddProbeToPort 向指定端口添加探测器
func (pm *ProbeMapping) AddProbeToPort(port int, probe string) {
	pm.mu.Lock()
	defer pm.mu.Unlock()

	if probe == "" {
		return
	}

	if existing, exists := pm.portMap[port]; exists {
		// 检查是否已存在
		for _, p := range existing {
			if p == probe {
				return
			}
		}
		// 添加到现有列表
		pm.portMap[port] = append(existing, probe)
	} else {
		// 创建新列表
		pm.portMap[port] = []string{probe}
	}
}

// RemoveProbeFromPort 从指定端口移除探测器
func (pm *ProbeMapping) RemoveProbeFromPort(port int, probe string) {
	pm.mu.Lock()
	defer pm.mu.Unlock()

	if existing, exists := pm.portMap[port]; exists {
		var result []string
		for _, p := range existing {
			if p != probe {
				result = append(result, p)
			}
		}

		if len(result) > 0 {
			pm.portMap[port] = result
		} else {
			delete(pm.portMap, port)
		}
	}
}

// GetDefaultProbes 获取默认探测器列表
func (pm *ProbeMapping) GetDefaultProbes() []string {
	pm.mu.RLock()
	defer pm.mu.RUnlock()

	result := make([]string, len(pm.defaultMap))
	copy(result, pm.defaultMap)
	return result
}

// SetDefaultProbes 设置默认探测器列表
func (pm *ProbeMapping) SetDefaultProbes(probes []string) {
	pm.mu.Lock()
	defer pm.mu.Unlock()

	if len(probes) > 0 {
		defaultCopy := make([]string, len(probes))
		copy(defaultCopy, probes)
		pm.defaultMap = defaultCopy
	}
}

// GetAllPortMappings 获取所有端口映射
func (pm *ProbeMapping) GetAllPortMappings() map[int][]string {
	pm.mu.RLock()
	defer pm.mu.RUnlock()

	result := make(map[int][]string)
	for port, probes := range pm.portMap {
		probesCopy := make([]string, len(probes))
		copy(probesCopy, probes)
		result[port] = probesCopy
	}
	return result
}

// GetMappedPorts 获取所有已映射的端口列表
func (pm *ProbeMapping) GetMappedPorts() []int {
	pm.mu.RLock()
	defer pm.mu.RUnlock()

	ports := make([]int, 0, len(pm.portMap))
	for port := range pm.portMap {
		ports = append(ports, port)
	}

	sort.Ints(ports)
	return ports
}

// HasMappingForPort 检查指定端口是否有自定义映射
func (pm *ProbeMapping) HasMappingForPort(port int) bool {
	pm.mu.RLock()
	defer pm.mu.RUnlock()

	_, exists := pm.portMap[port]
	return exists
}

// GetProbeCount 获取指定端口的探测器数量
func (pm *ProbeMapping) GetProbeCount(port int) int {
	pm.mu.RLock()
	defer pm.mu.RUnlock()

	if probes, exists := pm.portMap[port]; exists {
		return len(probes)
	}
	return len(pm.defaultMap)
}

// ResetToDefaults 重置为默认映射
func (pm *ProbeMapping) ResetToDefaults() {
	pm.mu.Lock()
	defer pm.mu.Unlock()

	pm.defaultMap = getDefaultProbeMap()
	pm.portMap = getDefaultPortMap()
}

// ClearPortMapping 清除指定端口的映射
func (pm *ProbeMapping) ClearPortMapping(port int) {
	pm.mu.Lock()
	defer pm.mu.Unlock()

	delete(pm.portMap, port)
}

// ClearAllMappings 清除所有端口映射
func (pm *ProbeMapping) ClearAllMappings() {
	pm.mu.Lock()
	defer pm.mu.Unlock()

	pm.portMap = make(map[int][]string)
}

// 全局探测器映射实例
var (
	globalProbeMapping *ProbeMapping
	probeMappingOnce   sync.Once
)

// GetGlobalProbeMapping 获取全局探测器映射实例
func GetGlobalProbeMapping() *ProbeMapping {
	probeMappingOnce.Do(func() {
		globalProbeMapping = NewProbeMapping()
	})
	return globalProbeMapping
}