network-topology-discovery/internal/topology/builder.go

package topology

import (
	"fmt"
	"network-topology-discovery/pkg/models"
	"strings"
)

// Builder 拓扑构建器
type Builder struct {
	devices []models.Device
}

// NewBuilder 创建拓扑构建器
func NewBuilder() *Builder {
	return &Builder{}
}

// AddDevice 添加设备（按IP去重）
func (b *Builder) AddDevice(device models.Device) {
	for i, existing := range b.devices {
		if existing.IP == device.IP {
			b.devices[i] = device // 覆盖更新
			return
		}
	}
	b.devices = append(b.devices, device)
}

// AddDevices 批量添加设备
func (b *Builder) AddDevices(devices []models.Device) {
	b.devices = append(b.devices, devices...)
}

// RemoveDevice 根据ID或IP移除设备
func (b *Builder) RemoveDevice(id string) {
	for i, d := range b.devices {
		if d.ID == id || d.IP == id {
			b.devices = append(b.devices[:i], b.devices[i+1:]...)
			return
		}
	}
}

// Build 构建拓扑图
func (b *Builder) Build() models.TopologyGraph {
	graph := models.TopologyGraph{
		Nodes: make([]models.TopologyNode, 0),
		Edges: make([]models.TopologyEdge, 0),
	}

	// 构建节点
	nodeMap := make(map[string]models.TopologyNode)
	for _, device := range b.devices {
		// 优先使用主机名,如果没有则使用IP作为显示名
		hostname := device.Hostname
		if hostname == "" {
			hostname = device.IP
		}

		node := models.TopologyNode{
			ID:       device.IP,
			IP:       device.IP,
			Hostname: hostname,
			Type:     string(device.Type),
			Icon:     getDeviceIcon(device.Type),
		}
		nodeMap[device.IP] = node
		graph.Nodes = append(graph.Nodes, node)
	}

	fmt.Printf("\n========== TOPOLOGY BUILD ===========\n")
	fmt.Printf("Total devices: %d\n", len(b.devices))
	fmt.Printf("Total nodes created: %d\n", len(graph.Nodes))
	for _, node := range graph.Nodes {
		fmt.Printf("  Node: %s (%s) - Type: %s\n", node.ID, node.Hostname, node.Type)
	}

	// 构建边(基于邻居信息)
	edgeMap := make(map[string]bool) // 用于去重

	// 打印所有设备的MAC地址，用于调试
	fmt.Println("\n========== MAC Address Database ==========")
	for _, d := range b.devices {
		if len(d.MACAddresses) > 0 {
			fmt.Printf("Device %s (%s): %d MAC addresses\n", d.Hostname, d.IP, len(d.MACAddresses))
			for i, mac := range d.MACAddresses {
				if i < 5 { // 只打印前5个
					fmt.Printf("  MAC[%d]: %s\n", i, mac)
				}
			}
			if len(d.MACAddresses) > 5 {
				fmt.Printf("  ... and %d more\n", len(d.MACAddresses)-5)
			}
		}
	}
	fmt.Println("==========================================\n")

	for _, device := range b.devices {
		fmt.Printf("Building edges for device %s: %d neighbors, %d MAC addresses\n",
			device.IP, len(device.Neighbors), len(device.MACAddresses))

		for _, neighbor := range device.Neighbors {
			// 过滤虚拟机邻居：识别MAC地址前缀，只保留真实网络设备
			if neighbor.RemoteMAC != "" {
				macPrefix := getMACPrefix(neighbor.RemoteMAC)
				if isVirtualMAC(macPrefix) {
					fmt.Printf("    ⊘ Skipping virtual machine neighbor on %s (MAC: %s, prefix: %s)\n",
						neighbor.LocalInterface, neighbor.RemoteMAC, macPrefix)
					continue
				}
			}

			fmt.Printf("  Processing neighbor on %s: RemoteMAC=%s, RemoteDevice=%s, RemoteIP=%s\n",
				neighbor.LocalInterface, neighbor.RemoteMAC, neighbor.RemoteDevice, neighbor.RemoteIP)

			// 尝试匹配邻居设备 - 多重策略
			targetIP := neighbor.RemoteIP
			matchMethod := "IP"

			// 策略1: 如果已有IP，直接使用
			if targetIP != "" {
				fmt.Printf("    ✓ Using direct IP: %s\n", targetIP)
			} else if neighbor.RemoteDevice != "" {
				// 策略2: 尝试通过主机名匹配
				fmt.Printf("    Trying hostname match: %s\n", neighbor.RemoteDevice)
				for _, d := range b.devices {
					if d.Hostname == neighbor.RemoteDevice {
						targetIP = d.IP
						matchMethod = "hostname"
						fmt.Printf("    ✓ Matched by hostname: %s -> %s\n", d.Hostname, d.IP)
						break
					}
				}
			}

			// 策略3: 通过MAC地址匹配
			if targetIP == "" && neighbor.RemoteMAC != "" {
				fmt.Printf("    Trying MAC match: %s\n", neighbor.RemoteMAC)
				for _, d := range b.devices {
					for _, mac := range d.MACAddresses {
						// 标准化MAC地址进行比较（去除分隔符，转小写）
						normalizedNeighborMAC := normalizeMAC(neighbor.RemoteMAC)
						normalizedDeviceMAC := normalizeMAC(mac)

						if normalizedNeighborMAC == normalizedDeviceMAC {
							targetIP = d.IP
							matchMethod = "MAC"
							fmt.Printf("    ✓ Matched by MAC: %s (device) == %s (neighbor) -> %s\n",
								mac, neighbor.RemoteMAC, d.IP)
							break
						}
					}
					if targetIP != "" {
						break
					}
				}
			}

			// 策略3b: MAC前缀匹配（使用12字符前缀 - 更精确）
			// 同品牌设备虽然4字符前缀相同，但12字符前缀通常不同
			if targetIP == "" && neighbor.RemoteMAC != "" {
				// 使用12字符前缀（例如：642f-c7e0-03）
				macPrefix := ""
				if len(neighbor.RemoteMAC) >= 12 {
					macPrefix = neighbor.RemoteMAC[:12]
				} else {
					macPrefix = neighbor.RemoteMAC
				}

				fmt.Printf("    Trying MAC prefix match (12 chars): %s (prefix: %s)\n", neighbor.RemoteMAC, macPrefix)

				var matchedDevices []string

				for _, d := range b.devices {
					if d.IP == device.IP {
						continue // 跳过自己
					}

					// 检查设备d的MAC地址是否有匹配的前缀
					matchingMACs := 0
					for _, mac := range d.MACAddresses {
						if len(mac) >= 12 && mac[:12] == macPrefix {
							matchingMACs++
						}
					}

					if matchingMACs >= 3 {
						matchedDevices = append(matchedDevices, d.IP)
					}
				}

				// 只有唯一匹配时才建立连接
				if len(matchedDevices) == 1 {
					targetIP = matchedDevices[0]
					matchMethod = fmt.Sprintf("MAC-prefix-12(%s)", macPrefix)
					fmt.Printf("    ✓ Matched by MAC prefix-12 (unique): %s -> %s\n", neighbor.RemoteMAC, targetIP)
				} else if len(matchedDevices) > 1 {
					fmt.Printf("    ✗ Ambiguous MAC prefix-12 match: %s matched %v\n", macPrefix, matchedDevices)
				} else {
					fmt.Printf("    ✗ No device matches MAC prefix-12 %s\n", macPrefix)
				}
			}

			// 策略3c: LLDP接口匹配（改进 - 单向但要求唯一性）
			// 如果设备A的LLDP信息显示它连接到接口Y
			// 而设备B有接口Y，则A连接到B
			// 但如果有多个设备都有接口Y，则不匹配（避免歧义）
			if targetIP == "" && neighbor.RemoteMAC != "" && neighbor.RemoteInterface != "" {
				fmt.Printf("    Trying LLDP interface match: looking for device with interface %s\n", neighbor.RemoteInterface)

				var matchedDevices []string

				for _, d := range b.devices {
					if d.IP == device.IP {
						continue // 跳过自己
					}

					// 检查设备d是否有这个接口
					for _, iface := range d.Interfaces {
						if iface.Name == neighbor.RemoteInterface {
							matchedDevices = append(matchedDevices, d.IP)
							break
						}
					}
				}

				// 只有唯一匹配时才建立连接
				if len(matchedDevices) == 1 {
					targetIP = matchedDevices[0]
					matchMethod = fmt.Sprintf("LLDP-interface(%s->%s)", neighbor.LocalInterface, neighbor.RemoteInterface)
					fmt.Printf("    ✓ Matched by LLDP interface (unique): %s %s -> %s %s\n",
						device.IP, neighbor.LocalInterface, targetIP, neighbor.RemoteInterface)
				} else if len(matchedDevices) > 1 {
					fmt.Printf("    ✗ Ambiguous: multiple devices have interface %s: %v\n",
						neighbor.RemoteInterface, matchedDevices)
				} else {
					fmt.Printf("    ✗ No device found with interface %s\n", neighbor.RemoteInterface)
				}
			}

			// 策略4: 通过本地接口IP网段匹配（新增）
			// 注意：此策略容易产生误匹配，仅在必要时使用
			if targetIP == "" {
				// 查找本地接口的IP地址
				localInterfaceIP := ""
				for _, iface := range device.Interfaces {
					if iface.Name == neighbor.LocalInterface && iface.IP != "" {
						localInterfaceIP = iface.IP
						break
					}
				}

				if localInterfaceIP != "" {
					fmt.Printf("    ⚠ Trying subnet match (risky): local interface %s has IP %s\n",
						neighbor.LocalInterface, localInterfaceIP)

					// 计算本地接口的网段
					localSubnet := getSubnet(localInterfaceIP)

					// 查找其他设备中是否有接口在同一网段
					for _, d := range b.devices {
						if d.IP == device.IP {
							continue // 跳过自己
						}

						// 检查设备的管理IP是否在同一网段
						if getSubnet(d.IP) == localSubnet {
							targetIP = d.IP
							matchMethod = "subnet(management IP)"
							fmt.Printf("    ⚠ Matched by subnet (risky): %s in %s\n", d.IP, localSubnet)
							break
						}

						// 检查设备的所有接口是否在同一网段
						for _, iface := range d.Interfaces {
							if iface.IP != "" && getSubnet(iface.IP) == localSubnet {
								targetIP = d.IP
								matchMethod = fmt.Sprintf("subnet(interface %s)", iface.Name)
								fmt.Printf("    ⚠ Matched by subnet (risky): %s (%s) in %s\n",
									d.IP, iface.Name, localSubnet)
								break
							}
						}
						if targetIP != "" {
							break
						}
					}
				} else {
					// 策略4b: 本地接口没有IP，尝试使用设备管理IP进行子网匹配（高风险策略 - 已禁用）
					// 此策略容易产生错误连接，暂时禁用
					fmt.Printf("    ✗ Skipping subnet match with management IP (disabled - too risky)\n")
				}
			}

			// 策略5: 通过接口描述匹配（新增）
			if targetIP == "" {
				// 查找本地接口的描述，看是否包含其他设备的IP或主机名
				for _, iface := range device.Interfaces {
					if iface.Name == neighbor.LocalInterface && iface.Description != "" {
						desc := strings.ToLower(iface.Description)
						fmt.Printf("    Trying description match: %s\n", iface.Description)

						for _, d := range b.devices {
							if d.IP == device.IP {
								continue
							}

							// 检查描述中是否包含目标设备的IP或主机名
							if strings.Contains(desc, strings.ToLower(d.IP)) ||
								(d.Hostname != "" && strings.Contains(desc, strings.ToLower(d.Hostname))) {
								targetIP = d.IP
								matchMethod = "description"
								fmt.Printf("    ✓ Matched by description: %s\n", iface.Description)
								break
							}
						}
					}
					if targetIP != "" {
						break
					}
				}
			}

			// 策略6: 通过邻居端口号反向匹配（已禁用 - 容易产生错误连接）
			// 暂时禁用，因为端口号匹配不够精确，可能产生环路

			if targetIP == "" {
				fmt.Printf("    ✗ Could not match neighbor on %s (MAC: %s)\n", neighbor.LocalInterface, neighbor.RemoteMAC)
				continue
			}

			fmt.Printf("  ✓ Creating edge: %s (%s) -> %s via %s, matched by %s\n",
				device.IP, neighbor.LocalInterface, targetIP, neighbor.RemoteInterface, matchMethod)

			// 避免重复边：同一对设备之间只保留一条边（不区分接口）
			// 如果这对设备已经有边了，跳过
			hasExistingEdge := false
			for _, existingEdge := range graph.Edges {
				if (existingEdge.Source == device.IP && existingEdge.Target == targetIP) ||
					(existingEdge.Source == targetIP && existingEdge.Target == device.IP) {
					fmt.Printf("    ⚠ Skipping: edge between %s and %s already exists\n", device.IP, targetIP)
					hasExistingEdge = true
					break
				}
			}
			if hasExistingEdge {
				continue
			}

			// 创建唯一的边ID（包含接口信息用于调试）
			edgeID := fmt.Sprintf("%s-%s-%s", device.IP, neighbor.LocalInterface, targetIP)

			edge := models.TopologyEdge{
				ID:              edgeID,
				Source:          device.IP,
				Target:          targetIP,
				SourceInterface: neighbor.LocalInterface,
				TargetInterface: neighbor.RemoteInterface,
				Protocol:        neighbor.Protocol,
			}

			graph.Edges = append(graph.Edges, edge)
			edgeMap[edgeID] = true
		}
	}

	fmt.Printf("\nTotal edges created: %d\n", len(graph.Edges))
	for _, edge := range graph.Edges {
		fmt.Printf("  Edge: %s (%s) -> %s (%s) [%s]\n",
			edge.Source, edge.SourceInterface,
			edge.Target, edge.TargetInterface,
			edge.Protocol)
	}
	fmt.Println("========================================\n")

	return graph
}

// normalizeMAC 标准化MAC地址（去除分隔符，转小写）
func normalizeMAC(mac string) string {
	// 去除所有分隔符（-、:、.）
	result := ""
	for _, c := range mac {
		if c != '-' && c != ':' && c != '.' {
			result += string(c)
		}
	}
	// 转小写
	return strings.ToLower(result)
}

// getMACPrefix 获取MAC地址前缀（前4个字符）
func getMACPrefix(mac string) string {
	normalized := normalizeMAC(mac)
	if len(normalized) >= 4 {
		return normalized[0:4]
	}
	return normalized
}

// isVirtualMAC 判断是否是虚拟机MAC地址前缀
func isVirtualMAC(prefix string) bool {
	// 根据用户反馈，这些前缀都是虚拟机
	userVirtualPrefixes := []string{"9844", "9400", "cc05", "44a1"}

	for _, vp := range userVirtualPrefixes {
		if prefix == vp {
			return true
		}
	}

	// 常见虚拟机/虚拟网卡MAC前缀（VMware、Hyper-V、KVM等）
	virtualPrefixes := []string{
		"0005", "000c", "0015", "0050", // VMware
		"0017", "001d", "0025", "5254", // KVM/QEMU
		"0015", "0017", "00ff", // Xen
	}

	for _, vp := range virtualPrefixes {
		if prefix == vp {
			return true
		}
	}

	return false
}

// getSubnet 获取IP地址的C类网段（简化版，适用于/24网络）
func getSubnet(ip string) string {
	parts := strings.Split(ip, ".")
	if len(parts) == 4 {
		// 返回前3段作为网段，例如: 172.16.8.0/24 -> "172.16.8"
		return fmt.Sprintf("%s.%s.%s", parts[0], parts[1], parts[2])
	}
	return ip
}

// getDeviceIcon 获取设备图标
func getDeviceIcon(deviceType models.DeviceType) string {
	switch deviceType {
	case models.DeviceTypeCisco:
		return "router"
	case models.DeviceTypeHuawei:
		return "router"
	case models.DeviceTypeH3C:
		return "switch"
	case models.DeviceTypeASA:
		return "firewall"
	case models.DeviceTypeLinux:
		return "server"
	case models.DeviceTypeWindows:
		return "server"
	default:
		return "device"
	}
}

// GetDevices 获取所有设备
func (b *Builder) GetDevices() []models.Device {
	return b.devices
}

// Clear 清空拓扑
func (b *Builder) Clear() {
	b.devices = make([]models.Device, 0)
}