network-topology-discovery/cmd/main.go

package main

import (
	"crypto/rand"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"sync"
	"time"

	"network-topology-discovery/internal/config"
	"network-topology-discovery/internal/device"
	"network-topology-discovery/internal/scanner"
	"network-topology-discovery/internal/storage"
	"network-topology-discovery/internal/terminal"
	"network-topology-discovery/internal/topology"
	"network-topology-discovery/pkg/models"
)

// App 应用
type App struct {
	config          *config.Config
	builder         *topology.Builder
	storage         *storage.Storage
	topologyStorage *storage.TopologyStorage
	tasks           map[string]*models.ScanTask
	mu              sync.RWMutex
	httpServer      *http.Server
}

// NewApp 创建应用
func NewApp(cfg *config.Config) *App {
	// 初始化拓扑存储（管理多个拓扑）
	topoStorage, err := storage.NewTopologyStorage("data")
	if err != nil {
		log.Printf("Warning: failed to initialize topology storage: %v", err)
	}

	// 初始化设备存储（使用默认文件，兼容旧版）
	store, err := storage.NewStorage("devices.json")
	if err != nil {
		log.Printf("Warning: failed to initialize storage: %v", err)
	}

	app := &App{
		config:          cfg,
		builder:         topology.NewBuilder(),
		storage:         store,
		topologyStorage: topoStorage,
		tasks:           make(map[string]*models.ScanTask),
	}

	// 如果有拓扑存储，切换到第一个拓扑
	if topoStorage != nil {
		topos, err := topoStorage.GetAllTopologies()
		if err == nil && len(topos) > 0 {
			topoStorage.SetCurrentTopology(topos[0].ID)
			deviceFile := topoStorage.GetDeviceFilePath(topos[0].ID)
			app.storage.SetFilePath(deviceFile)
			log.Printf("Loaded topology: %s", topos[0].Name)
		}
	}

	// 加载设备到拓扑构建器
	if store != nil {
		devices, err := store.GetAllDevices()
		if err != nil {
			log.Printf("Warning: failed to load devices from database: %v", err)
		} else {
			log.Printf("Loaded %d devices from storage", len(devices))
			for _, dev := range devices {
				app.builder.AddDevice(dev)
			}
		}
	}

	return app
}

// Start 启动应用
func (app *App) Start() error {
	// 设置路由
	mux := http.NewServeMux()

	// 静态文件服务 - 使用文件系统而非embed
	webDir := getWebDir()
	if _, err := os.Stat(webDir); err == nil {
		mux.Handle("/", http.FileServer(http.Dir(webDir)))
	} else {
		log.Printf("警告: web目录不存在,静态文件服务不可用")
		mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
			w.Write([]byte("<h1>网络拓扑发现系统</h1><p>Web界面文件未找到</p>"))
		})
	}

	// API路由
	mux.HandleFunc("/api/scan", app.handleScan)
	mux.HandleFunc("/api/scan/{id}", app.handleScanProgress)
	mux.HandleFunc("/api/topology", app.handleTopology)
	mux.HandleFunc("/api/devices", app.handleGetDevices)
	mux.HandleFunc("/api/device", app.handleAddDevice)
	mux.HandleFunc("/api/device/{id}", app.handleDeviceDetail)

	// 拓扑管理API
	mux.HandleFunc("/api/topologies", app.handleTopologies)
	mux.HandleFunc("/api/topology/switch", app.handleSwitchTopology)
	mux.HandleFunc("/api/topology/{id}", app.handleTopologyDetail)

	// SSH终端API
	mux.HandleFunc("/api/terminal", app.handleTerminalConnect)

	// 全局设备池API
	mux.HandleFunc("/api/devices/all", app.handleGetAllDevices)
	mux.HandleFunc("/api/topology/{id}/devices", app.handleAddDevicesToTopology)

	addr := fmt.Sprintf("%s:%d", app.config.Web.Host, app.config.Web.Port)
	app.httpServer = &http.Server{
		Addr:    addr,
		Handler: mux,
	}

	log.Printf("服务启动在 %s", addr)
	return app.httpServer.ListenAndServe()
}

// 生成唯一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:])
}

// getWebDir 获取web目录路径
func getWebDir() string {
	// 尝试多个可能的路径
	possiblePaths := []string{
		"web",
		filepath.Join("cmd", "web"),
		filepath.Join("..", "web"),
	}

	for _, path := range possiblePaths {
		if _, err := os.Stat(path); err == nil {
			absPath, _ := filepath.Abs(path)
			return absPath
		}
	}

	// 默认返回web
	return "web"
}

// 处理扫描请求
func (app *App) handleScan(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	var req struct {
		ScanRange string `json:"scan_range"`
		SSHPort   int    `json:"ssh_port"`
		Username  string `json:"username"`
		Password  string `json:"password"`
	}

	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	if req.SSHPort == 0 {
		req.SSHPort = 22
	}

	// 创建扫描任务
	taskID := generateID()
	task := &models.ScanTask{
		ID:        taskID,
		Status:    "running",
		StartTime: time.Now(),
		Devices:   []models.Device{},
	}

	app.mu.Lock()
	app.tasks[taskID] = task
	app.mu.Unlock()

	// 确保当前有拓扑
	if app.topologyStorage != nil && app.topologyStorage.GetCurrentTopologyID() == "" {
		topos, err := app.topologyStorage.GetAllTopologies()
		if err == nil && len(topos) > 0 {
			app.topologyStorage.SetCurrentTopology(topos[0].ID)
			deviceFile := app.topologyStorage.GetDeviceFilePath(topos[0].ID)
			app.storage.SetFilePath(deviceFile)
		}
	}

	// 异步执行扫描
	go app.runScan(task, req.ScanRange, req.SSHPort, req.Username, req.Password)

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]string{"task_id": taskID})
}

// 执行扫描
func (app *App) runScan(task *models.ScanTask, cidr string, sshPort int, username, password string) {
	defer func() {
		task.EndTime = time.Now()
	}()

	// 创建扫描器
	sc := scanner.NewScanner(app.config.Scanner.Concurrency, time.Duration(app.config.Scanner.Timeout)*time.Second)

	// 阶段1: 解析IP范围 (进度 0% -> 10%)
	ips, err := sc.ScanRange(cidr)
	if err != nil {
		task.Status = "failed"
		task.ErrorMessage = err.Error()
		return
	}
	task.TotalDevices = len(ips)
	task.Progress = 10
	log.Printf("[扫描] 阶段1: 解析CIDR %s, 共 %d 个IP", cidr, len(ips))

	// 阶段2: 检查存活主机 (进度 10% -> 30%)
	aliveHosts := sc.CheckHosts(ips)
	task.Progress = 30
	log.Printf("[扫描] 阶段2: 发现 %d 个存活主机", len(aliveHosts))

	// 阶段3: 检查SSH端口 (进度 30% -> 50%)
	sshHosts := sc.CheckSSHHosts(aliveHosts, sshPort)
	task.Progress = 50
	task.TotalDevices = len(sshHosts)
	log.Printf("[扫描] 阶段3: 发现 %d 个SSH主机", len(sshHosts))

	// 如果没有SSH主机，直接完成
	if len(sshHosts) == 0 {
		task.Status = "completed"
		task.Progress = 100
		log.Printf("[扫描] 未发现SSH主机，扫描完成")
		return
	}

	// 阶段4: 采集设备信息 (进度 50% -> 100%)
	var devices []models.Device
	for i, ip := range sshHosts {
		log.Printf("[扫描] 阶段4: 正在采集设备 %s (%d/%d)", ip, i+1, len(sshHosts))

		// 尝试不同设备类型
		deviceTypes := []models.DeviceType{
			models.DeviceTypeCisco,
			models.DeviceTypeHuawei,
			models.DeviceTypeH3C,
			models.DeviceTypeASA,
			models.DeviceTypeLinux,
			models.DeviceTypeWindows,
		}

		var discoveredDevice *models.Device
		for _, dtype := range deviceTypes {
			dev, err := device.DiscoverDevice(ip, dtype, username, password)
			if err == nil && dev.ScanStatus == "success" {
				discoveredDevice = dev
				break
			}
		}

		if discoveredDevice != nil {
			devices = append(devices, *discoveredDevice)
			app.builder.AddDevice(*discoveredDevice)

			// 保存到数据库
			if app.storage != nil {
				if err := app.storage.SaveDevice(discoveredDevice); err != nil {
					log.Printf("Warning: failed to save device %s to database: %v", ip, err)
				}
			}
		}

		// 更新进度: 50% ~ 100%
		task.ScannedDevices = i + 1
		task.Progress = 50 + (i+1)*50/len(sshHosts)
		task.Devices = devices
	}

	task.Status = "completed"
	task.Progress = 100
	task.Devices = devices
	log.Printf("[扫描] 完成，共发现 %d 台设备", len(devices))
}

// 处理扫描进度查询
func (app *App) handleScanProgress(w http.ResponseWriter, r *http.Request) {
	id := r.PathValue("id")

	app.mu.RLock()
	task, exists := app.tasks[id]
	app.mu.RUnlock()

	if !exists {
		http.Error(w, "Task not found", http.StatusNotFound)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(task)
}

// 处理拓扑查询
func (app *App) handleTopology(w http.ResponseWriter, r *http.Request) {
	graph := app.builder.Build()

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(graph)
}

// 处理获取所有设备
func (app *App) handleGetDevices(w http.ResponseWriter, r *http.Request) {
	var devices []models.Device

	// 优先从存储获取
	if app.storage != nil {
		var err error
		devices, err = app.storage.GetAllDevices()
		if err != nil {
			log.Printf("Error: failed to get devices from storage: %v", err)
			// 降级到 builder获取
			devices = app.builder.GetDevices()
		}
		log.Printf("Returning %d devices from storage", len(devices))
	} else {
		devices = app.builder.GetDevices()
		log.Printf("Returning %d devices from builder", len(devices))
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(devices)
}

// 处理添加设备
func (app *App) handleAddDevice(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	var req struct {
		IP       string `json:"ip"`
		Type     string `json:"type"`
		Username string `json:"username"`
		Password string `json:"password"`
	}

	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	deviceType := models.DeviceType(req.Type)
	log.Printf("Adding device: %s (type: %s)", req.IP, req.Type)

	dev, err := device.DiscoverDevice(req.IP, deviceType, req.Username, req.Password)
	if err != nil {
		log.Printf("Failed to discover device %s: %v", req.IP, err)
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusInternalServerError)
		json.NewEncoder(w).Encode(map[string]string{"message": err.Error()})
		return
	}

	log.Printf("Device discovered: %s, interfaces: %d, neighbors: %d",
		dev.IP, len(dev.Interfaces), len(dev.Neighbors))

	app.builder.AddDevice(*dev)

	// 保存到存储
	if app.storage != nil {
		if err := app.storage.SaveDevice(dev); err != nil {
			log.Printf("Error: failed to save device %s to storage: %v", req.IP, err)
		} else {
			log.Printf("Device %s saved to storage successfully", req.IP)
		}
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(dev)
}

// 处理设备详情查询
func (app *App) handleDeviceDetail(w http.ResponseWriter, r *http.Request) {
	id := r.PathValue("id")

	devices := app.builder.GetDevices()
	for _, dev := range devices {
		if dev.ID == id || dev.IP == id {
			w.Header().Set("Content-Type", "application/json")
			json.NewEncoder(w).Encode(dev)
			return
		}
	}

	http.Error(w, "Device not found", http.StatusNotFound)
}

// 处理拓扑列表（GET: 获取所有拓扑，POST: 创建新拓扑）
func (app *App) handleTopologies(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")

	if r.Method == http.MethodGet {
		// 获取所有拓扑
		if app.topologyStorage == nil {
			json.NewEncoder(w).Encode([]models.Topology{})
			return
		}

		topos, err := app.topologyStorage.GetAllTopologies()
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		// 更新每个拓扑的设备数量
		currentTopoID := app.topologyStorage.GetCurrentTopologyID()
		for i := range topos {
			deviceFile := app.topologyStorage.GetDeviceFilePath(topos[i].ID)
			store, err := storage.NewStorage(deviceFile)
			if err == nil {
				devices, err := store.GetAllDevices()
				if err == nil {
					topos[i].DeviceCount = len(devices)
				}
			}
			// 标记当前拓扑
			if topos[i].ID == currentTopoID {
				topos[i].Name = topos[i].Name + " (当前)"
			}
		}

		json.NewEncoder(w).Encode(topos)
	} else if r.Method == http.MethodPost {
		// 创建新拓扑
		var req struct {
			Name        string `json:"name"`
			Description string `json:"description"`
			ScanRange   string `json:"scan_range"`
			SSHPort     int    `json:"ssh_port"`
			Username    string `json:"username"`
		}

		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		if req.Name == "" {
			http.Error(w, "name is required", http.StatusBadRequest)
			return
		}

		topo, err := app.topologyStorage.CreateTopology(req.Name, req.Description, req.ScanRange, req.SSHPort, req.Username)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		// 自动切换到新拓扑
		app.topologyStorage.SetCurrentTopology(topo.ID)
		deviceFile := app.topologyStorage.GetDeviceFilePath(topo.ID)
		app.storage.SetFilePath(deviceFile)
		app.builder.Clear()

		log.Printf("Created and switched to new topology: %s", topo.Name)
		json.NewEncoder(w).Encode(topo)
	} else {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
	}
}

// 处理切换拓扑
func (app *App) handleSwitchTopology(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	var req struct {
		TopologyID string `json:"topology_id"`
	}

	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	if app.topologyStorage == nil {
		http.Error(w, "Topology storage not available", http.StatusInternalServerError)
		return
	}

	// 验证拓扑存在
	topo, err := app.topologyStorage.GetTopology(req.TopologyID)
	if err != nil {
		http.Error(w, "Topology not found", http.StatusNotFound)
		return
	}

	// 切换拓扑
	err = app.topologyStorage.SetCurrentTopology(req.TopologyID)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// 切换设备存储文件
	deviceFile := app.topologyStorage.GetDeviceFilePath(req.TopologyID)
	err = app.storage.SetFilePath(deviceFile)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// 清空并重新加载拓扑构建器
	app.builder.Clear()
	devices, err := app.storage.GetAllDevices()
	if err == nil {
		for _, dev := range devices {
			app.builder.AddDevice(dev)
		}
	}

	log.Printf("Switched to topology: %s (%s)", topo.Name, req.TopologyID)
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]string{"message": "Topology switched successfully"})
}

// 处理拓扑详情
func (app *App) handleTopologyDetail(w http.ResponseWriter, r *http.Request) {
	id := r.PathValue("id")

	if app.topologyStorage == nil {
		http.Error(w, "Topology storage not available", http.StatusInternalServerError)
		return
	}

	topo, err := app.topologyStorage.GetTopology(id)
	if err != nil {
		http.Error(w, "Topology not found", http.StatusNotFound)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(topo)
}

// 处理SSH终端连接
func (app *App) handleTerminalConnect(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		// WebSocket 连接使用 GET，但我们通过 POST 获取凭据后再升级
		if r.Method != http.MethodGet {
			http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
			return
		}
	}

	// 从查询参数获取设备IP
	ip := r.URL.Query().Get("ip")
	if ip == "" {
		http.Error(w, "ip parameter is required", http.StatusBadRequest)
		return
	}

	port := 22
	if p := r.URL.Query().Get("port"); p != "" {
		fmt.Sscanf(p, "%d", &port)
	}

	username := r.URL.Query().Get("username")
	password := r.URL.Query().Get("password")

	if username == "" || password == "" {
		http.Error(w, "username and password are required", http.StatusBadRequest)
		return
	}

	log.Printf("[终端] 正在连接 %s:%d (%s)", ip, port, username)

	// 使用 terminal handler 处理 WebSocket 连接
	terminal.HandleTerminal(w, r, ip, port, username, password)
}

// 获取所有拓扑中的全部设备（全局设备池）
func (app *App) handleGetAllDevices(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	deviceMap := make(map[string]models.Device) // 用IP去重

	if app.topologyStorage != nil {
		topos, err := app.topologyStorage.GetAllTopologies()
		if err == nil {
			for _, topo := range topos {
				deviceFile := app.topologyStorage.GetDeviceFilePath(topo.ID)
				store, err := storage.NewStorage(deviceFile)
				if err != nil {
					continue
				}
				devices, err := store.GetAllDevices()
				if err != nil {
					continue
				}
				for _, dev := range devices {
					deviceMap[dev.IP] = dev
				}
			}
		}
	}

	// 也从旧版 devices.json 加载
	if app.storage != nil {
		devices, err := app.storage.GetAllDevices()
		if err == nil {
			for _, dev := range devices {
				deviceMap[dev.IP] = dev
			}
		}
	}

	devices := make([]models.Device, 0, len(deviceMap))
	for _, dev := range deviceMap {
		devices = append(devices, dev)
	}

	log.Printf("[设备池] 返回 %d 个全局设备", len(devices))
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(devices)
}

// 向指定拓扑批量添加设备
func (app *App) handleAddDevicesToTopology(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	topoID := r.PathValue("id")

	var req struct {
		DeviceIDs []string `json:"device_ids"` // IP列表或ID列表
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	if len(req.DeviceIDs) == 0 {
		http.Error(w, "device_ids is required", http.StatusBadRequest)
		return
	}

	// 获取目标拓扑的存储
	if app.topologyStorage == nil {
		http.Error(w, "Topology storage not available", http.StatusInternalServerError)
		return
	}

	topo, err := app.topologyStorage.GetTopology(topoID)
	if err != nil {
		http.Error(w, "Topology not found", http.StatusNotFound)
		return
	}

	deviceFile := app.topologyStorage.GetDeviceFilePath(topoID)
	targetStore, err := storage.NewStorage(deviceFile)
	if err != nil {
		http.Error(w, "Failed to open target storage", http.StatusInternalServerError)
		return
	}

	// 从全局设备池中查找并添加
	allDeviceMap := make(map[string]models.Device)
	topos, _ := app.topologyStorage.GetAllTopologies()
	for _, t := range topos {
		df := app.topologyStorage.GetDeviceFilePath(t.ID)
		s, err := storage.NewStorage(df)
		if err != nil {
			continue
		}
		devs, err := s.GetAllDevices()
		if err != nil {
			continue
		}
		for _, d := range devs {
			allDeviceMap[d.IP] = d
			allDeviceMap[d.ID] = d
		}
	}

	added := 0
	for _, idOrIP := range req.DeviceIDs {
		if dev, ok := allDeviceMap[idOrIP]; ok {
			devCopy := dev // 复制一份
			targetStore.SaveDevice(&devCopy)
			added++
		}
	}

	// 如果是当前拓扑，刷新builder
	if app.topologyStorage.GetCurrentTopologyID() == topoID {
		app.builder.Clear()
		devices, _ := targetStore.GetAllDevices()
		for _, dev := range devices {
			app.builder.AddDevice(dev)
		}
	}

	log.Printf("[拓扑] 向 %s 添加了 %d/%d 设备", topo.Name, added, len(req.DeviceIDs))
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]interface{}{
		"added":    added,
		"total":    len(req.DeviceIDs),
		"topology": topo.Name,
	})
}

func main() {
	// 加载配置
	configFile := "config.json"
	if len(os.Args) > 1 {
		configFile = os.Args[1]
	}

	var cfg *config.Config
	if _, err := os.Stat(configFile); err == nil {
		cfg, err = config.LoadConfig(configFile)
		if err != nil {
			log.Printf("加载配置文件失败: %v, 使用默认配置", err)
			cfg = config.DefaultConfig()
		}
	} else {
		log.Printf("配置文件不存在, 使用默认配置")
		cfg = config.DefaultConfig()
	}

	// 创建并启动应用
	app := NewApp(cfg)

	log.Println("网络拓扑发现系统启动...")
	if err := app.Start(); err != nil {
		log.Fatalf("服务启动失败: %v", err)
	}
}