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

package topology

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

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

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

// AddDevice 添加设备
func (b *Builder) AddDevice(device models.Device) {
	b.devices = append(b.devices, device)
}

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

// 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)
	}

	// 构建边(基于邻居信息)
	edgeMap := make(map[string]bool) // 用于去重
	for _, device := range b.devices {
		for _, neighbor := range device.Neighbors {
			// 尝试匹配邻居设备
			targetIP := neighbor.RemoteIP
			
			// 如果没有IP，尝试通过设备名匹配
			if targetIP == "" && neighbor.RemoteDevice != "" {
				for _, d := range b.devices {
					if d.Hostname == neighbor.RemoteDevice {
						targetIP = d.IP
						break
					}
				}
			}
			
			// 如果还是没有，尝试通过MAC地址匹配（新增）
			if targetIP == "" && neighbor.RemoteMAC != "" {
				for _, d := range b.devices {
					for _, mac := range d.MACAddresses {
						if mac == neighbor.RemoteMAC {
							targetIP = d.IP
							break
						}
					}
					if targetIP != "" {
						break
					}
				}
			}

			if targetIP == "" {
				continue
			}

			// 创建唯一的边ID
			edgeID := fmt.Sprintf("%s-%s-%s", device.IP, neighbor.LocalInterface, targetIP)
			reverseEdgeID := fmt.Sprintf("%s-%s-%s", targetIP, neighbor.RemoteInterface, device.IP)

			// 避免重复边
			if edgeMap[edgeID] || edgeMap[reverseEdgeID] {
				continue
			}

			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
		}
	}

	return graph
}

// 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)
}