prod

2026-04-26 22:00:03 +08:00
parent bcad5b8e27
commit 2a97f458a9
11 changed files with 888 additions and 53 deletions
@@ -8,6 +8,7 @@ import (
 	"net/http"
 	"os"
 	"path/filepath"
+	"strings"
 	"sync"
 	"time"

@@ -29,6 +30,8 @@ type App struct {
 	tasks           map[string]*models.ScanTask
 	mu              sync.RWMutex
 	httpServer      *http.Server
+	sessions        map[string]time.Time // token -> expire time
+	sessionMu       sync.RWMutex
 }

 // NewApp 创建应用
@@ -51,6 +54,7 @@ func NewApp(cfg *config.Config) *App {
 		storage:         store,
 		topologyStorage: topoStorage,
 		tasks:           make(map[string]*models.ScanTask),
+		sessions:        make(map[string]time.Time),
 	}

 	// 如果有拓扑存储，切换到第一个拓扑
@@ -85,36 +89,56 @@ func (app *App) Start() error {
 	// 设置路由
 	mux := http.NewServeMux()

-	// 静态文件服务 - 使用文件系统而非embed
+	// 登录API（无需认证）
+	mux.HandleFunc("/api/login", app.handleLogin)
+	mux.HandleFunc("/api/logout", app.handleLogout)
+
+	// 登录页面（无需认证）
+	mux.HandleFunc("/login", app.handleLoginPage)
+
+	// 认证中间件保护的路由
+	authMux := http.NewServeMux()
+
+	// 静态文件服务
 	webDir := getWebDir()
 	if _, err := os.Stat(webDir); err == nil {
-		mux.Handle("/", http.FileServer(http.Dir(webDir)))
+		authMux.Handle("/", http.FileServer(http.Dir(webDir)))
 	} else {
 		log.Printf("警告: web目录不存在,静态文件服务不可用")
-		mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+		authMux.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)
+	authMux.HandleFunc("/api/scan", app.handleScan)
+	authMux.HandleFunc("/api/scan/{id}", app.handleScanProgress)
+	authMux.HandleFunc("/api/topology", app.handleTopology)
+	authMux.HandleFunc("/api/devices", app.handleGetDevices)
+	authMux.HandleFunc("/api/device", app.handleAddDevice)
+	authMux.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)
+	authMux.HandleFunc("/api/topologies", app.handleTopologies)
+	authMux.HandleFunc("/api/topology/switch", app.handleSwitchTopology)
+	authMux.HandleFunc("/api/topology/{id}", app.handleTopologyDetail)

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

 	// 全局设备池API
-	mux.HandleFunc("/api/devices/all", app.handleGetAllDevices)
-	mux.HandleFunc("/api/topology/{id}/devices", app.handleAddDevicesToTopology)
+	authMux.HandleFunc("/api/devices/all", app.handleGetAllDevices)
+	authMux.HandleFunc("/api/topology/{id}/devices", app.handleAddDevicesToTopology)
+
+	// 根据认证配置决定是否启用中间件
+	var handler http.Handler = authMux
+	if app.config.Auth.Enabled {
+		handler = app.authMiddleware(authMux)
+		log.Printf("认证已启用, 用户: %s", app.config.Auth.Username)
+	} else {
+		log.Printf("认证未启用")
+	}
+	mux.Handle("/", handler)

 	addr := fmt.Sprintf("%s:%d", app.config.Web.Host, app.config.Web.Port)
 	app.httpServer = &http.Server{
@@ -394,6 +418,64 @@ func (app *App) handleAddDevice(w http.ResponseWriter, r *http.Request) {
 func (app *App) handleDeviceDetail(w http.ResponseWriter, r *http.Request) {
 	id := r.PathValue("id")

+	if r.Method == http.MethodDelete {
+		// 删除设备
+		if app.storage == nil {
+			http.Error(w, "Storage not available", http.StatusInternalServerError)
+			return
+		}
+		if err := app.storage.DeleteDevice(id); err != nil {
+			log.Printf("Failed to delete device %s: %v", id, err)
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+		// 同时从 builder 中移除
+		app.builder.RemoveDevice(id)
+		log.Printf("Deleted device: %s", id)
+		w.Header().Set("Content-Type", "application/json")
+		json.NewEncoder(w).Encode(map[string]string{"status": "success"})
+		return
+	}
+
+	if r.Method == http.MethodPut {
+		// 修改设备信息（Hostname、Type）
+		var req struct {
+			Hostname string `json:"hostname"`
+			Type     string `json:"type"`
+		}
+		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+			http.Error(w, err.Error(), http.StatusBadRequest)
+			return
+		}
+
+		if app.storage == nil {
+			http.Error(w, "Storage not available", http.StatusInternalServerError)
+			return
+		}
+		dev, err := app.storage.GetDevice(id)
+		if err != nil {
+			http.Error(w, "Device not found", http.StatusNotFound)
+			return
+		}
+		if req.Hostname != "" {
+			dev.Hostname = req.Hostname
+		}
+		if req.Type != "" {
+			dev.Type = models.DeviceType(req.Type)
+		}
+		if err := app.storage.SaveDevice(dev); err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+		// 同时更新 builder 中的设备
+		app.builder.AddDevice(*dev)
+		log.Printf("Updated device: %s", id)
+		w.Header().Set("Content-Type", "application/json")
+		json.NewEncoder(w).Encode(dev)
+		return
+	}
+
+	// GET: 获取设备详情
 	devices := app.builder.GetDevices()
 	for _, dev := range devices {
 		if dev.ID == id || dev.IP == id {
@@ -592,6 +674,125 @@ func (app *App) handleTerminalConnect(w http.ResponseWriter, r *http.Request) {
 	terminal.HandleTerminal(w, r, ip, port, username, password)
 }

+// ==================== 认证相关 ====================
+
+// authMiddleware 认证中间件
+func (app *App) authMiddleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// 从 cookie 获取 token
+		cookie, err := r.Cookie("session_token")
+		if err != nil || !app.isValidSession(cookie.Value) {
+			// API 请求返回 401，页面请求重定向到登录
+			if strings.HasPrefix(r.URL.Path, "/api/") {
+				http.Error(w, `{"error": "Unauthorized"}`, http.StatusUnauthorized)
+			} else {
+				http.Redirect(w, r, "/login", http.StatusFound)
+			}
+			return
+		}
+		next.ServeHTTP(w, r)
+	})
+}
+
+// isValidSession 验证会话是否有效
+func (app *App) isValidSession(token string) bool {
+	if token == "" {
+		return false
+	}
+	app.sessionMu.RLock()
+	expire, exists := app.sessions[token]
+	app.sessionMu.RUnlock()
+	if !exists {
+		return false
+	}
+	if time.Now().After(expire) {
+		app.sessionMu.Lock()
+		delete(app.sessions, token)
+		app.sessionMu.Unlock()
+		return false
+	}
+	return true
+}
+
+// generateToken 生成随机 token
+func generateToken() string {
+	b := make([]byte, 32)
+	rand.Read(b)
+	return fmt.Sprintf("%x", b)
+}
+
+// handleLogin 处理登录
+func (app *App) handleLogin(w http.ResponseWriter, r *http.Request) {
+	if r.Method != http.MethodPost {
+		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
+		return
+	}
+
+	var req struct {
+		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 !app.config.Auth.Enabled {
+		json.NewEncoder(w).Encode(map[string]interface{}{"success": true})
+		return
+	}
+
+	if req.Username == app.config.Auth.Username && req.Password == app.config.Auth.Password {
+		token := generateToken()
+		expire := time.Now().Add(24 * time.Hour)
+
+		app.sessionMu.Lock()
+		app.sessions[token] = expire
+		app.sessionMu.Unlock()
+
+		http.SetCookie(w, &http.Cookie{
+			Name:     "session_token",
+			Value:    token,
+			Path:     "/",
+			Expires:  expire,
+			HttpOnly: true,
+		})
+
+		log.Printf("用户 %s 登录成功", req.Username)
+		json.NewEncoder(w).Encode(map[string]interface{}{"success": true})
+	} else {
+		log.Printf("登录失败: 用户名或密码错误 (username=%s)", req.Username)
+		w.WriteHeader(http.StatusUnauthorized)
+		json.NewEncoder(w).Encode(map[string]interface{}{"success": false, "message": "用户名或密码错误"})
+	}
+}
+
+// handleLogout 处理登出
+func (app *App) handleLogout(w http.ResponseWriter, r *http.Request) {
+	if cookie, err := r.Cookie("session_token"); err == nil {
+		app.sessionMu.Lock()
+		delete(app.sessions, cookie.Value)
+		app.sessionMu.Unlock()
+	}
+	http.SetCookie(w, &http.Cookie{
+		Name:     "session_token",
+		Value:    "",
+		Path:     "/",
+		MaxAge:   -1,
+		HttpOnly: true,
+	})
+	http.Redirect(w, r, "/login", http.StatusFound)
+}
+
+// handleLoginPage 返回登录页面
+func (app *App) handleLoginPage(w http.ResponseWriter, r *http.Request) {
+	if !app.config.Auth.Enabled {
+		http.Redirect(w, r, "/", http.StatusFound)
+		return
+	}
+	http.ServeFile(w, r, filepath.Join(getWebDir(), "login.html"))
+}
+
 // 获取所有拓扑中的全部设备（全局设备池）
 func (app *App) handleGetAllDevices(w http.ResponseWriter, r *http.Request) {
 	if r.Method != http.MethodGet {