package storage import ( "crypto/rand" "encoding/json" "fmt" "log" "os" "path/filepath" "sync" "time" "network-topology-discovery/pkg/models" ) // TopologyStorage 拓扑存储管理 type TopologyStorage struct { mu sync.RWMutex dataDir string topologies map[string]*models.Topology currentTopoID string } // NewTopologyStorage 创建拓扑存储 func NewTopologyStorage(dataDir string) (*TopologyStorage, error) { s := &TopologyStorage{ dataDir: dataDir, topologies: make(map[string]*models.Topology), } // 确保数据目录存在 if err := os.MkdirAll(dataDir, 0755); err != nil { return nil, fmt.Errorf("failed to create data directory: %w", err) } // 加载拓扑元数据 if err := s.loadMeta(); err != nil { if !os.IsNotExist(err) { return nil, fmt.Errorf("failed to load topology meta: %w", err) } log.Printf("Creating new topology storage at %s", dataDir) } return s, nil } // loadMeta 加载拓扑元数据 func (s *TopologyStorage) loadMeta() error { metaFile := filepath.Join(s.dataDir, "topologies.json") data, err := os.ReadFile(metaFile) if err != nil { return err } var topos []models.Topology if err := json.Unmarshal(data, &topos); err != nil { return fmt.Errorf("failed to parse topology meta: %w", err) } for i := range topos { s.topologies[topos[i].ID] = &topos[i] } log.Printf("Loaded %d topologies from meta", len(topos)) return nil } // saveMeta 保存拓扑元数据 func (s *TopologyStorage) saveMeta() error { topos := make([]models.Topology, 0, len(s.topologies)) for _, topo := range s.topologies { topos = append(topos, *topo) } data, err := json.MarshalIndent(topos, "", " ") if err != nil { return fmt.Errorf("failed to marshal topologies: %w", err) } metaFile := filepath.Join(s.dataDir, "topologies.json") if err := os.WriteFile(metaFile, data, 0644); err != nil { return fmt.Errorf("failed to write topology meta: %w", err) } return nil } // CreateTopology 创建新拓扑 func (s *TopologyStorage) CreateTopology(name, description, scanRange string, sshPort int, username string) (*models.Topology, error) { s.mu.Lock() defer s.mu.Unlock() topo := &models.Topology{ ID: generateID(), Name: name, Description: description, ScanRange: scanRange, SSHPort: sshPort, Username: username, CreatedAt: time.Now(), UpdatedAt: time.Now(), DeviceCount: 0, } s.topologies[topo.ID] = topo // 创建该拓扑的设备文件 deviceFile := filepath.Join(s.dataDir, topo.ID+"_devices.json") if _, err := os.Stat(deviceFile); os.IsNotExist(err) { if err := os.WriteFile(deviceFile, []byte("[]"), 0644); err != nil { return nil, fmt.Errorf("failed to create device file: %w", err) } } // 保存元数据 if err := s.saveMeta(); err != nil { return nil, err } log.Printf("Topology created: %s (%s)", topo.Name, topo.ID) return topo, nil } // GetTopology 获取拓扑 func (s *TopologyStorage) GetTopology(id string) (*models.Topology, error) { s.mu.RLock() defer s.mu.RUnlock() topo, exists := s.topologies[id] if !exists { return nil, fmt.Errorf("topology not found: %s", id) } return topo, nil } // GetAllTopologies 获取所有拓扑 func (s *TopologyStorage) GetAllTopologies() ([]models.Topology, error) { s.mu.RLock() defer s.mu.RUnlock() topos := make([]models.Topology, 0, len(s.topologies)) for _, topo := range s.topologies { topos = append(topos, *topo) } return topos, nil } // UpdateTopology 更新拓扑 func (s *TopologyStorage) UpdateTopology(id string, updates map[string]interface{}) error { s.mu.Lock() defer s.mu.Unlock() topo, exists := s.topologies[id] if !exists { return fmt.Errorf("topology not found: %s", id) } if name, ok := updates["name"].(string); ok { topo.Name = name } if desc, ok := updates["description"].(string); ok { topo.Description = desc } if scanRange, ok := updates["scan_range"].(string); ok { topo.ScanRange = scanRange } if sshPort, ok := updates["ssh_port"].(float64); ok { topo.SSHPort = int(sshPort) } if username, ok := updates["username"].(string); ok { topo.Username = username } topo.UpdatedAt = time.Now() return s.saveMeta() } // DeleteTopology 删除拓扑及其所有数据 func (s *TopologyStorage) DeleteTopology(id string) error { s.mu.Lock() defer s.mu.Unlock() if _, exists := s.topologies[id]; !exists { return fmt.Errorf("topology not found: %s", id) } // 删除设备文件 deviceFile := filepath.Join(s.dataDir, id+"_devices.json") if err := os.Remove(deviceFile); err != nil && !os.IsNotExist(err) { log.Printf("Warning: failed to delete device file for topology %s: %v", id, err) } delete(s.topologies, id) // 更新当前拓扑 if s.currentTopoID == id { s.currentTopoID = "" } return s.saveMeta() } // SetCurrentTopology 设置当前拓扑 func (s *TopologyStorage) SetCurrentTopology(id string) error { if _, err := s.GetTopology(id); err != nil { return err } s.mu.Lock() defer s.mu.Unlock() s.currentTopoID = id log.Printf("Current topology set to: %s", id) return nil } // GetCurrentTopologyID 获取当前拓扑ID func (s *TopologyStorage) GetCurrentTopologyID() string { s.mu.RLock() defer s.mu.RUnlock() return s.currentTopoID } // GetDeviceFilePath 获取拓扑的设备文件路径 func (s *TopologyStorage) GetDeviceFilePath(topoID string) string { return filepath.Join(s.dataDir, topoID+"_devices.json") } // generateID 生成唯一ID func generateID() string { b := make([]byte, 16) rand.Read(b) return fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:]) }