main.go 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958
  1. package main
  2. import (
  3. "crypto/rand"
  4. "encoding/json"
  5. "fmt"
  6. "log"
  7. "net/http"
  8. "os"
  9. "path/filepath"
  10. "strings"
  11. "sync"
  12. "time"
  13. "network-topology-discovery/internal/config"
  14. "network-topology-discovery/internal/device"
  15. "network-topology-discovery/internal/scanner"
  16. "network-topology-discovery/internal/storage"
  17. "network-topology-discovery/internal/terminal"
  18. "network-topology-discovery/internal/topology"
  19. "network-topology-discovery/pkg/models"
  20. )
  21. // App 应用
  22. type App struct {
  23. config *config.Config
  24. builder *topology.Builder
  25. storage *storage.Storage
  26. topologyStorage *storage.TopologyStorage
  27. tasks map[string]*models.ScanTask
  28. mu sync.RWMutex
  29. httpServer *http.Server
  30. sessions map[string]time.Time // token -> expire time
  31. sessionMu sync.RWMutex
  32. }
  33. // NewApp 创建应用
  34. func NewApp(cfg *config.Config) *App {
  35. // 初始化拓扑存储(管理多个拓扑)
  36. topoStorage, err := storage.NewTopologyStorage("data")
  37. if err != nil {
  38. log.Printf("Warning: failed to initialize topology storage: %v", err)
  39. }
  40. // 初始化设备存储(使用默认文件,兼容旧版)
  41. store, err := storage.NewStorage("devices.json")
  42. if err != nil {
  43. log.Printf("Warning: failed to initialize storage: %v", err)
  44. }
  45. app := &App{
  46. config: cfg,
  47. builder: topology.NewBuilder(),
  48. storage: store,
  49. topologyStorage: topoStorage,
  50. tasks: make(map[string]*models.ScanTask),
  51. sessions: make(map[string]time.Time),
  52. }
  53. // 如果有拓扑存储,切换到第一个拓扑
  54. if topoStorage != nil {
  55. topos, err := topoStorage.GetAllTopologies()
  56. if err == nil && len(topos) > 0 {
  57. topoStorage.SetCurrentTopology(topos[0].ID)
  58. deviceFile := topoStorage.GetDeviceFilePath(topos[0].ID)
  59. app.storage.SetFilePath(deviceFile)
  60. log.Printf("Loaded topology: %s", topos[0].Name)
  61. }
  62. }
  63. // 加载设备到拓扑构建器
  64. if store != nil {
  65. devices, err := store.GetAllDevices()
  66. if err != nil {
  67. log.Printf("Warning: failed to load devices from database: %v", err)
  68. } else {
  69. log.Printf("Loaded %d devices from storage", len(devices))
  70. for _, dev := range devices {
  71. app.builder.AddDevice(dev)
  72. }
  73. }
  74. }
  75. return app
  76. }
  77. // Start 启动应用
  78. func (app *App) Start() error {
  79. // 设置路由
  80. mux := http.NewServeMux()
  81. // 登录API(无需认证)
  82. mux.HandleFunc("/api/login", app.handleLogin)
  83. mux.HandleFunc("/api/logout", app.handleLogout)
  84. // 登录页面(无需认证)
  85. mux.HandleFunc("/login", app.handleLoginPage)
  86. // 认证中间件保护的路由
  87. authMux := http.NewServeMux()
  88. // 静态文件服务
  89. webDir := getWebDir()
  90. if _, err := os.Stat(webDir); err == nil {
  91. authMux.Handle("/", http.FileServer(http.Dir(webDir)))
  92. } else {
  93. log.Printf("警告: web目录不存在,静态文件服务不可用")
  94. authMux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
  95. w.Write([]byte("<h1>网络拓扑发现系统</h1><p>Web界面文件未找到</p>"))
  96. })
  97. }
  98. // API路由
  99. authMux.HandleFunc("/api/scan", app.handleScan)
  100. authMux.HandleFunc("/api/scan/{id}", app.handleScanProgress)
  101. authMux.HandleFunc("/api/topology", app.handleTopology)
  102. authMux.HandleFunc("/api/devices", app.handleGetDevices)
  103. authMux.HandleFunc("/api/device", app.handleAddDevice)
  104. authMux.HandleFunc("/api/device/{id}", app.handleDeviceDetail)
  105. // 拓扑管理API
  106. authMux.HandleFunc("/api/topologies", app.handleTopologies)
  107. authMux.HandleFunc("/api/topology/switch", app.handleSwitchTopology)
  108. authMux.HandleFunc("/api/topology/{id}", app.handleTopologyDetail)
  109. // SSH终端API
  110. authMux.HandleFunc("/api/terminal", app.handleTerminalConnect)
  111. // 全局设备池API
  112. authMux.HandleFunc("/api/devices/all", app.handleGetAllDevices)
  113. authMux.HandleFunc("/api/topology/{id}/devices", app.handleAddDevicesToTopology)
  114. // 根据认证配置决定是否启用中间件
  115. var handler http.Handler = authMux
  116. if app.config.Auth.Enabled {
  117. handler = app.authMiddleware(authMux)
  118. log.Printf("认证已启用, 用户: %s", app.config.Auth.Username)
  119. } else {
  120. log.Printf("认证未启用")
  121. }
  122. mux.Handle("/", handler)
  123. addr := fmt.Sprintf("%s:%d", app.config.Web.Host, app.config.Web.Port)
  124. app.httpServer = &http.Server{
  125. Addr: addr,
  126. Handler: mux,
  127. }
  128. log.Printf("服务启动在 %s", addr)
  129. return app.httpServer.ListenAndServe()
  130. }
  131. // 生成唯一ID
  132. func generateID() string {
  133. b := make([]byte, 16)
  134. rand.Read(b)
  135. return fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:])
  136. }
  137. // getWebDir 获取web目录路径
  138. func getWebDir() string {
  139. // 尝试多个可能的路径
  140. possiblePaths := []string{
  141. "web",
  142. filepath.Join("cmd", "web"),
  143. filepath.Join("..", "web"),
  144. }
  145. for _, path := range possiblePaths {
  146. if _, err := os.Stat(path); err == nil {
  147. absPath, _ := filepath.Abs(path)
  148. return absPath
  149. }
  150. }
  151. // 默认返回web
  152. return "web"
  153. }
  154. // 处理扫描请求
  155. func (app *App) handleScan(w http.ResponseWriter, r *http.Request) {
  156. if r.Method != http.MethodPost {
  157. http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
  158. return
  159. }
  160. var req struct {
  161. ScanRange string `json:"scan_range"`
  162. SSHPort int `json:"ssh_port"`
  163. Username string `json:"username"`
  164. Password string `json:"password"`
  165. }
  166. if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
  167. http.Error(w, err.Error(), http.StatusBadRequest)
  168. return
  169. }
  170. if req.SSHPort == 0 {
  171. req.SSHPort = 22
  172. }
  173. // 创建扫描任务
  174. taskID := generateID()
  175. task := &models.ScanTask{
  176. ID: taskID,
  177. Status: "running",
  178. StartTime: time.Now(),
  179. Devices: []models.Device{},
  180. }
  181. app.mu.Lock()
  182. app.tasks[taskID] = task
  183. app.mu.Unlock()
  184. // 确保当前有拓扑
  185. if app.topologyStorage != nil && app.topologyStorage.GetCurrentTopologyID() == "" {
  186. topos, err := app.topologyStorage.GetAllTopologies()
  187. if err == nil && len(topos) > 0 {
  188. app.topologyStorage.SetCurrentTopology(topos[0].ID)
  189. deviceFile := app.topologyStorage.GetDeviceFilePath(topos[0].ID)
  190. app.storage.SetFilePath(deviceFile)
  191. }
  192. }
  193. // 异步执行扫描
  194. go app.runScan(task, req.ScanRange, req.SSHPort, req.Username, req.Password)
  195. w.Header().Set("Content-Type", "application/json")
  196. json.NewEncoder(w).Encode(map[string]string{"task_id": taskID})
  197. }
  198. // 执行扫描
  199. func (app *App) runScan(task *models.ScanTask, cidr string, sshPort int, username, password string) {
  200. defer func() {
  201. task.EndTime = time.Now()
  202. }()
  203. // 创建扫描器
  204. sc := scanner.NewScanner(app.config.Scanner.Concurrency, time.Duration(app.config.Scanner.Timeout)*time.Second)
  205. // 阶段1: 解析IP范围 (进度 0% -> 10%)
  206. ips, err := sc.ScanRange(cidr)
  207. if err != nil {
  208. task.Status = "failed"
  209. task.ErrorMessage = err.Error()
  210. return
  211. }
  212. task.TotalDevices = len(ips)
  213. task.Progress = 10
  214. log.Printf("[扫描] 阶段1: 解析CIDR %s, 共 %d 个IP", cidr, len(ips))
  215. // 阶段2: 检查存活主机 (进度 10% -> 30%)
  216. aliveHosts := sc.CheckHosts(ips)
  217. task.Progress = 30
  218. log.Printf("[扫描] 阶段2: 发现 %d 个存活主机", len(aliveHosts))
  219. // 阶段3: 检查SSH端口 (进度 30% -> 50%)
  220. sshHosts := sc.CheckSSHHosts(aliveHosts, sshPort)
  221. task.Progress = 50
  222. task.TotalDevices = len(sshHosts)
  223. log.Printf("[扫描] 阶段3: 发现 %d 个SSH主机", len(sshHosts))
  224. // 如果没有SSH主机,直接完成
  225. if len(sshHosts) == 0 {
  226. task.Status = "completed"
  227. task.Progress = 100
  228. log.Printf("[扫描] 未发现SSH主机,扫描完成")
  229. return
  230. }
  231. // 阶段4: 采集设备信息 (进度 50% -> 100%)
  232. var devices []models.Device
  233. for i, ip := range sshHosts {
  234. log.Printf("[扫描] 阶段4: 正在采集设备 %s (%d/%d)", ip, i+1, len(sshHosts))
  235. // 尝试不同设备类型
  236. deviceTypes := []models.DeviceType{
  237. models.DeviceTypeCisco,
  238. models.DeviceTypeHuawei,
  239. models.DeviceTypeH3C,
  240. models.DeviceTypeASA,
  241. models.DeviceTypeLinux,
  242. models.DeviceTypeWindows,
  243. }
  244. var discoveredDevice *models.Device
  245. for _, dtype := range deviceTypes {
  246. dev, err := device.DiscoverDevice(ip, dtype, username, password)
  247. if err == nil && dev.ScanStatus == "success" {
  248. discoveredDevice = dev
  249. break
  250. }
  251. }
  252. if discoveredDevice != nil {
  253. devices = append(devices, *discoveredDevice)
  254. app.builder.AddDevice(*discoveredDevice)
  255. // 保存到数据库
  256. if app.storage != nil {
  257. if err := app.storage.SaveDevice(discoveredDevice); err != nil {
  258. log.Printf("Warning: failed to save device %s to database: %v", ip, err)
  259. }
  260. }
  261. }
  262. // 更新进度: 50% ~ 100%
  263. task.ScannedDevices = i + 1
  264. task.Progress = 50 + (i+1)*50/len(sshHosts)
  265. task.Devices = devices
  266. }
  267. task.Status = "completed"
  268. task.Progress = 100
  269. task.Devices = devices
  270. log.Printf("[扫描] 完成,共发现 %d 台设备", len(devices))
  271. }
  272. // 处理扫描进度查询
  273. func (app *App) handleScanProgress(w http.ResponseWriter, r *http.Request) {
  274. id := r.PathValue("id")
  275. app.mu.RLock()
  276. task, exists := app.tasks[id]
  277. app.mu.RUnlock()
  278. if !exists {
  279. http.Error(w, "Task not found", http.StatusNotFound)
  280. return
  281. }
  282. w.Header().Set("Content-Type", "application/json")
  283. json.NewEncoder(w).Encode(task)
  284. }
  285. // 处理拓扑查询
  286. func (app *App) handleTopology(w http.ResponseWriter, r *http.Request) {
  287. graph := app.builder.Build()
  288. w.Header().Set("Content-Type", "application/json")
  289. json.NewEncoder(w).Encode(graph)
  290. }
  291. // 处理获取所有设备
  292. func (app *App) handleGetDevices(w http.ResponseWriter, r *http.Request) {
  293. var devices []models.Device
  294. // 优先从存储获取
  295. if app.storage != nil {
  296. var err error
  297. devices, err = app.storage.GetAllDevices()
  298. if err != nil {
  299. log.Printf("Error: failed to get devices from storage: %v", err)
  300. // 降级到 builder获取
  301. devices = app.builder.GetDevices()
  302. }
  303. log.Printf("Returning %d devices from storage", len(devices))
  304. } else {
  305. devices = app.builder.GetDevices()
  306. log.Printf("Returning %d devices from builder", len(devices))
  307. }
  308. w.Header().Set("Content-Type", "application/json")
  309. json.NewEncoder(w).Encode(devices)
  310. }
  311. // 处理添加设备
  312. func (app *App) handleAddDevice(w http.ResponseWriter, r *http.Request) {
  313. if r.Method != http.MethodPost {
  314. http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
  315. return
  316. }
  317. var req struct {
  318. IP string `json:"ip"`
  319. Type string `json:"type"`
  320. Username string `json:"username"`
  321. Password string `json:"password"`
  322. }
  323. if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
  324. http.Error(w, err.Error(), http.StatusBadRequest)
  325. return
  326. }
  327. deviceType := models.DeviceType(req.Type)
  328. log.Printf("Adding device: %s (type: %s)", req.IP, req.Type)
  329. dev, err := device.DiscoverDevice(req.IP, deviceType, req.Username, req.Password)
  330. if err != nil {
  331. log.Printf("Failed to discover device %s: %v", req.IP, err)
  332. w.Header().Set("Content-Type", "application/json")
  333. w.WriteHeader(http.StatusInternalServerError)
  334. json.NewEncoder(w).Encode(map[string]string{"message": err.Error()})
  335. return
  336. }
  337. log.Printf("Device discovered: %s, interfaces: %d, neighbors: %d",
  338. dev.IP, len(dev.Interfaces), len(dev.Neighbors))
  339. app.builder.AddDevice(*dev)
  340. // 保存到存储
  341. if app.storage != nil {
  342. if err := app.storage.SaveDevice(dev); err != nil {
  343. log.Printf("Error: failed to save device %s to storage: %v", req.IP, err)
  344. } else {
  345. log.Printf("Device %s saved to storage successfully", req.IP)
  346. }
  347. }
  348. w.Header().Set("Content-Type", "application/json")
  349. json.NewEncoder(w).Encode(dev)
  350. }
  351. // 处理设备详情查询
  352. func (app *App) handleDeviceDetail(w http.ResponseWriter, r *http.Request) {
  353. id := r.PathValue("id")
  354. if r.Method == http.MethodDelete {
  355. // 删除设备
  356. if app.storage == nil {
  357. http.Error(w, "Storage not available", http.StatusInternalServerError)
  358. return
  359. }
  360. if err := app.storage.DeleteDevice(id); err != nil {
  361. log.Printf("Failed to delete device %s: %v", id, err)
  362. http.Error(w, err.Error(), http.StatusInternalServerError)
  363. return
  364. }
  365. // 同时从 builder 中移除
  366. app.builder.RemoveDevice(id)
  367. log.Printf("Deleted device: %s", id)
  368. w.Header().Set("Content-Type", "application/json")
  369. json.NewEncoder(w).Encode(map[string]string{"status": "success"})
  370. return
  371. }
  372. if r.Method == http.MethodPut {
  373. // 修改设备信息(Hostname、Type)
  374. var req struct {
  375. Hostname string `json:"hostname"`
  376. Type string `json:"type"`
  377. }
  378. if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
  379. http.Error(w, err.Error(), http.StatusBadRequest)
  380. return
  381. }
  382. if app.storage == nil {
  383. http.Error(w, "Storage not available", http.StatusInternalServerError)
  384. return
  385. }
  386. dev, err := app.storage.GetDevice(id)
  387. if err != nil {
  388. http.Error(w, "Device not found", http.StatusNotFound)
  389. return
  390. }
  391. if req.Hostname != "" {
  392. dev.Hostname = req.Hostname
  393. }
  394. if req.Type != "" {
  395. dev.Type = models.DeviceType(req.Type)
  396. }
  397. if err := app.storage.SaveDevice(dev); err != nil {
  398. http.Error(w, err.Error(), http.StatusInternalServerError)
  399. return
  400. }
  401. // 同时更新 builder 中的设备
  402. app.builder.AddDevice(*dev)
  403. log.Printf("Updated device: %s", id)
  404. w.Header().Set("Content-Type", "application/json")
  405. json.NewEncoder(w).Encode(dev)
  406. return
  407. }
  408. // GET: 获取设备详情
  409. devices := app.builder.GetDevices()
  410. for _, dev := range devices {
  411. if dev.ID == id || dev.IP == id {
  412. w.Header().Set("Content-Type", "application/json")
  413. json.NewEncoder(w).Encode(dev)
  414. return
  415. }
  416. }
  417. http.Error(w, "Device not found", http.StatusNotFound)
  418. }
  419. // 处理拓扑列表(GET: 获取所有拓扑,POST: 创建新拓扑)
  420. func (app *App) handleTopologies(w http.ResponseWriter, r *http.Request) {
  421. w.Header().Set("Content-Type", "application/json")
  422. if r.Method == http.MethodGet {
  423. // 获取所有拓扑
  424. if app.topologyStorage == nil {
  425. json.NewEncoder(w).Encode([]models.Topology{})
  426. return
  427. }
  428. topos, err := app.topologyStorage.GetAllTopologies()
  429. if err != nil {
  430. http.Error(w, err.Error(), http.StatusInternalServerError)
  431. return
  432. }
  433. // 更新每个拓扑的设备数量
  434. currentTopoID := app.topologyStorage.GetCurrentTopologyID()
  435. for i := range topos {
  436. deviceFile := app.topologyStorage.GetDeviceFilePath(topos[i].ID)
  437. store, err := storage.NewStorage(deviceFile)
  438. if err == nil {
  439. devices, err := store.GetAllDevices()
  440. if err == nil {
  441. topos[i].DeviceCount = len(devices)
  442. }
  443. }
  444. // 标记当前拓扑
  445. if topos[i].ID == currentTopoID {
  446. topos[i].Name = topos[i].Name + " (当前)"
  447. }
  448. }
  449. json.NewEncoder(w).Encode(topos)
  450. } else if r.Method == http.MethodPost {
  451. // 创建新拓扑
  452. var req struct {
  453. Name string `json:"name"`
  454. Description string `json:"description"`
  455. ScanRange string `json:"scan_range"`
  456. SSHPort int `json:"ssh_port"`
  457. Username string `json:"username"`
  458. }
  459. if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
  460. http.Error(w, err.Error(), http.StatusBadRequest)
  461. return
  462. }
  463. if req.Name == "" {
  464. http.Error(w, "name is required", http.StatusBadRequest)
  465. return
  466. }
  467. topo, err := app.topologyStorage.CreateTopology(req.Name, req.Description, req.ScanRange, req.SSHPort, req.Username)
  468. if err != nil {
  469. http.Error(w, err.Error(), http.StatusInternalServerError)
  470. return
  471. }
  472. // 自动切换到新拓扑
  473. app.topologyStorage.SetCurrentTopology(topo.ID)
  474. deviceFile := app.topologyStorage.GetDeviceFilePath(topo.ID)
  475. app.storage.SetFilePath(deviceFile)
  476. app.builder.Clear()
  477. log.Printf("Created and switched to new topology: %s", topo.Name)
  478. json.NewEncoder(w).Encode(topo)
  479. } else {
  480. http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
  481. }
  482. }
  483. // 处理切换拓扑
  484. func (app *App) handleSwitchTopology(w http.ResponseWriter, r *http.Request) {
  485. if r.Method != http.MethodPost {
  486. http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
  487. return
  488. }
  489. var req struct {
  490. TopologyID string `json:"topology_id"`
  491. }
  492. if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
  493. http.Error(w, err.Error(), http.StatusBadRequest)
  494. return
  495. }
  496. if app.topologyStorage == nil {
  497. http.Error(w, "Topology storage not available", http.StatusInternalServerError)
  498. return
  499. }
  500. // 验证拓扑存在
  501. topo, err := app.topologyStorage.GetTopology(req.TopologyID)
  502. if err != nil {
  503. http.Error(w, "Topology not found", http.StatusNotFound)
  504. return
  505. }
  506. // 切换拓扑
  507. err = app.topologyStorage.SetCurrentTopology(req.TopologyID)
  508. if err != nil {
  509. http.Error(w, err.Error(), http.StatusInternalServerError)
  510. return
  511. }
  512. // 切换设备存储文件
  513. deviceFile := app.topologyStorage.GetDeviceFilePath(req.TopologyID)
  514. err = app.storage.SetFilePath(deviceFile)
  515. if err != nil {
  516. http.Error(w, err.Error(), http.StatusInternalServerError)
  517. return
  518. }
  519. // 清空并重新加载拓扑构建器
  520. app.builder.Clear()
  521. devices, err := app.storage.GetAllDevices()
  522. if err == nil {
  523. for _, dev := range devices {
  524. app.builder.AddDevice(dev)
  525. }
  526. }
  527. log.Printf("Switched to topology: %s (%s)", topo.Name, req.TopologyID)
  528. w.Header().Set("Content-Type", "application/json")
  529. json.NewEncoder(w).Encode(map[string]string{"message": "Topology switched successfully"})
  530. }
  531. // 处理拓扑详情
  532. func (app *App) handleTopologyDetail(w http.ResponseWriter, r *http.Request) {
  533. id := r.PathValue("id")
  534. if app.topologyStorage == nil {
  535. http.Error(w, "Topology storage not available", http.StatusInternalServerError)
  536. return
  537. }
  538. topo, err := app.topologyStorage.GetTopology(id)
  539. if err != nil {
  540. http.Error(w, "Topology not found", http.StatusNotFound)
  541. return
  542. }
  543. w.Header().Set("Content-Type", "application/json")
  544. json.NewEncoder(w).Encode(topo)
  545. }
  546. // 处理SSH终端连接
  547. func (app *App) handleTerminalConnect(w http.ResponseWriter, r *http.Request) {
  548. if r.Method != http.MethodPost {
  549. // WebSocket 连接使用 GET,但我们通过 POST 获取凭据后再升级
  550. if r.Method != http.MethodGet {
  551. http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
  552. return
  553. }
  554. }
  555. // 从查询参数获取设备IP
  556. ip := r.URL.Query().Get("ip")
  557. if ip == "" {
  558. http.Error(w, "ip parameter is required", http.StatusBadRequest)
  559. return
  560. }
  561. port := 22
  562. if p := r.URL.Query().Get("port"); p != "" {
  563. fmt.Sscanf(p, "%d", &port)
  564. }
  565. username := r.URL.Query().Get("username")
  566. password := r.URL.Query().Get("password")
  567. if username == "" || password == "" {
  568. http.Error(w, "username and password are required", http.StatusBadRequest)
  569. return
  570. }
  571. log.Printf("[终端] 正在连接 %s:%d (%s)", ip, port, username)
  572. // 使用 terminal handler 处理 WebSocket 连接
  573. terminal.HandleTerminal(w, r, ip, port, username, password)
  574. }
  575. // ==================== 认证相关 ====================
  576. // authMiddleware 认证中间件
  577. func (app *App) authMiddleware(next http.Handler) http.Handler {
  578. return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  579. // 从 cookie 获取 token
  580. cookie, err := r.Cookie("session_token")
  581. if err != nil || !app.isValidSession(cookie.Value) {
  582. // API 请求返回 401,页面请求重定向到登录
  583. if strings.HasPrefix(r.URL.Path, "/api/") {
  584. http.Error(w, `{"error": "Unauthorized"}`, http.StatusUnauthorized)
  585. } else {
  586. http.Redirect(w, r, "/login", http.StatusFound)
  587. }
  588. return
  589. }
  590. next.ServeHTTP(w, r)
  591. })
  592. }
  593. // isValidSession 验证会话是否有效
  594. func (app *App) isValidSession(token string) bool {
  595. if token == "" {
  596. return false
  597. }
  598. app.sessionMu.RLock()
  599. expire, exists := app.sessions[token]
  600. app.sessionMu.RUnlock()
  601. if !exists {
  602. return false
  603. }
  604. if time.Now().After(expire) {
  605. app.sessionMu.Lock()
  606. delete(app.sessions, token)
  607. app.sessionMu.Unlock()
  608. return false
  609. }
  610. return true
  611. }
  612. // generateToken 生成随机 token
  613. func generateToken() string {
  614. b := make([]byte, 32)
  615. rand.Read(b)
  616. return fmt.Sprintf("%x", b)
  617. }
  618. // handleLogin 处理登录
  619. func (app *App) handleLogin(w http.ResponseWriter, r *http.Request) {
  620. if r.Method != http.MethodPost {
  621. http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
  622. return
  623. }
  624. var req struct {
  625. Username string `json:"username"`
  626. Password string `json:"password"`
  627. }
  628. if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
  629. http.Error(w, err.Error(), http.StatusBadRequest)
  630. return
  631. }
  632. if !app.config.Auth.Enabled {
  633. json.NewEncoder(w).Encode(map[string]interface{}{"success": true})
  634. return
  635. }
  636. if req.Username == app.config.Auth.Username && req.Password == app.config.Auth.Password {
  637. token := generateToken()
  638. expire := time.Now().Add(24 * time.Hour)
  639. app.sessionMu.Lock()
  640. app.sessions[token] = expire
  641. app.sessionMu.Unlock()
  642. http.SetCookie(w, &http.Cookie{
  643. Name: "session_token",
  644. Value: token,
  645. Path: "/",
  646. Expires: expire,
  647. HttpOnly: true,
  648. })
  649. log.Printf("用户 %s 登录成功", req.Username)
  650. json.NewEncoder(w).Encode(map[string]interface{}{"success": true})
  651. } else {
  652. log.Printf("登录失败: 用户名或密码错误 (username=%s)", req.Username)
  653. w.WriteHeader(http.StatusUnauthorized)
  654. json.NewEncoder(w).Encode(map[string]interface{}{"success": false, "message": "用户名或密码错误"})
  655. }
  656. }
  657. // handleLogout 处理登出
  658. func (app *App) handleLogout(w http.ResponseWriter, r *http.Request) {
  659. if cookie, err := r.Cookie("session_token"); err == nil {
  660. app.sessionMu.Lock()
  661. delete(app.sessions, cookie.Value)
  662. app.sessionMu.Unlock()
  663. }
  664. http.SetCookie(w, &http.Cookie{
  665. Name: "session_token",
  666. Value: "",
  667. Path: "/",
  668. MaxAge: -1,
  669. HttpOnly: true,
  670. })
  671. http.Redirect(w, r, "/login", http.StatusFound)
  672. }
  673. // handleLoginPage 返回登录页面
  674. func (app *App) handleLoginPage(w http.ResponseWriter, r *http.Request) {
  675. if !app.config.Auth.Enabled {
  676. http.Redirect(w, r, "/", http.StatusFound)
  677. return
  678. }
  679. http.ServeFile(w, r, filepath.Join(getWebDir(), "login.html"))
  680. }
  681. // 获取所有拓扑中的全部设备(全局设备池)
  682. func (app *App) handleGetAllDevices(w http.ResponseWriter, r *http.Request) {
  683. if r.Method != http.MethodGet {
  684. http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
  685. return
  686. }
  687. deviceMap := make(map[string]models.Device) // 用IP去重
  688. if app.topologyStorage != nil {
  689. topos, err := app.topologyStorage.GetAllTopologies()
  690. if err == nil {
  691. for _, topo := range topos {
  692. deviceFile := app.topologyStorage.GetDeviceFilePath(topo.ID)
  693. store, err := storage.NewStorage(deviceFile)
  694. if err != nil {
  695. continue
  696. }
  697. devices, err := store.GetAllDevices()
  698. if err != nil {
  699. continue
  700. }
  701. for _, dev := range devices {
  702. deviceMap[dev.IP] = dev
  703. }
  704. }
  705. }
  706. }
  707. // 也从旧版 devices.json 加载
  708. if app.storage != nil {
  709. devices, err := app.storage.GetAllDevices()
  710. if err == nil {
  711. for _, dev := range devices {
  712. deviceMap[dev.IP] = dev
  713. }
  714. }
  715. }
  716. devices := make([]models.Device, 0, len(deviceMap))
  717. for _, dev := range deviceMap {
  718. devices = append(devices, dev)
  719. }
  720. log.Printf("[设备池] 返回 %d 个全局设备", len(devices))
  721. w.Header().Set("Content-Type", "application/json")
  722. json.NewEncoder(w).Encode(devices)
  723. }
  724. // 向指定拓扑批量添加设备
  725. func (app *App) handleAddDevicesToTopology(w http.ResponseWriter, r *http.Request) {
  726. if r.Method != http.MethodPost {
  727. http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
  728. return
  729. }
  730. topoID := r.PathValue("id")
  731. var req struct {
  732. DeviceIDs []string `json:"device_ids"` // IP列表或ID列表
  733. }
  734. if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
  735. http.Error(w, err.Error(), http.StatusBadRequest)
  736. return
  737. }
  738. if len(req.DeviceIDs) == 0 {
  739. http.Error(w, "device_ids is required", http.StatusBadRequest)
  740. return
  741. }
  742. // 获取目标拓扑的存储
  743. if app.topologyStorage == nil {
  744. http.Error(w, "Topology storage not available", http.StatusInternalServerError)
  745. return
  746. }
  747. topo, err := app.topologyStorage.GetTopology(topoID)
  748. if err != nil {
  749. http.Error(w, "Topology not found", http.StatusNotFound)
  750. return
  751. }
  752. deviceFile := app.topologyStorage.GetDeviceFilePath(topoID)
  753. targetStore, err := storage.NewStorage(deviceFile)
  754. if err != nil {
  755. http.Error(w, "Failed to open target storage", http.StatusInternalServerError)
  756. return
  757. }
  758. // 从全局设备池中查找并添加
  759. allDeviceMap := make(map[string]models.Device)
  760. topos, _ := app.topologyStorage.GetAllTopologies()
  761. for _, t := range topos {
  762. df := app.topologyStorage.GetDeviceFilePath(t.ID)
  763. s, err := storage.NewStorage(df)
  764. if err != nil {
  765. continue
  766. }
  767. devs, err := s.GetAllDevices()
  768. if err != nil {
  769. continue
  770. }
  771. for _, d := range devs {
  772. allDeviceMap[d.IP] = d
  773. allDeviceMap[d.ID] = d
  774. }
  775. }
  776. added := 0
  777. for _, idOrIP := range req.DeviceIDs {
  778. if dev, ok := allDeviceMap[idOrIP]; ok {
  779. devCopy := dev // 复制一份
  780. targetStore.SaveDevice(&devCopy)
  781. added++
  782. }
  783. }
  784. // 如果是当前拓扑,刷新builder
  785. if app.topologyStorage.GetCurrentTopologyID() == topoID {
  786. app.builder.Clear()
  787. devices, _ := targetStore.GetAllDevices()
  788. for _, dev := range devices {
  789. app.builder.AddDevice(dev)
  790. }
  791. }
  792. log.Printf("[拓扑] 向 %s 添加了 %d/%d 设备", topo.Name, added, len(req.DeviceIDs))
  793. w.Header().Set("Content-Type", "application/json")
  794. json.NewEncoder(w).Encode(map[string]interface{}{
  795. "added": added,
  796. "total": len(req.DeviceIDs),
  797. "topology": topo.Name,
  798. })
  799. }
  800. func main() {
  801. // 加载配置
  802. configFile := "config.json"
  803. if len(os.Args) > 1 {
  804. configFile = os.Args[1]
  805. }
  806. var cfg *config.Config
  807. if _, err := os.Stat(configFile); err == nil {
  808. cfg, err = config.LoadConfig(configFile)
  809. if err != nil {
  810. log.Printf("加载配置文件失败: %v, 使用默认配置", err)
  811. cfg = config.DefaultConfig()
  812. }
  813. } else {
  814. log.Printf("配置文件不存在, 使用默认配置")
  815. cfg = config.DefaultConfig()
  816. }
  817. // 创建并启动应用
  818. app := NewApp(cfg)
  819. log.Println("网络拓扑发现系统启动...")
  820. if err := app.Start(); err != nil {
  821. log.Fatalf("服务启动失败: %v", err)
  822. }
  823. }