chore: ARTEX 初始化

ARTEX — AI 自主渗透测试系统（Go 后端 + Next.js 前端，单二进制内嵌）。 Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
2026-07-26 04:27:45 -04:00
commit 0ef8561411
@@ -0,0 +1,224 @@
+package server
+
+import (
+	"crypto/rand"
+	"fmt"
+	"log"
+	"math/big"
+	"net/http"
+	"os"
+	"path/filepath"
+	"strings"
+	"time"
+
+	"github.com/golang-jwt/jwt/v5"
+	"golang.org/x/crypto/bcrypt"
+)
+
+const (
+	jwtKeyFilename = "jwt.key"
+	authPassKey    = "auth.password_hash"
+	jwtTTL         = 7 * 24 * time.Hour
+	keyChars       = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+)
+
+// loadOrCreateJWTKey reads the 32-byte signing key from dataDir/jwt.key.
+// On first run it generates a cryptographically random key and persists it.
+func loadOrCreateJWTKey(dataDir string) ([]byte, error) {
+	path := filepath.Join(dataDir, jwtKeyFilename)
+	if data, err := os.ReadFile(path); err == nil && len(strings.TrimSpace(string(data))) >= 32 {
+		return []byte(strings.TrimSpace(string(data))), nil
+	}
+	buf := make([]byte, 32)
+	for i := range buf {
+		n, err := rand.Int(rand.Reader, big.NewInt(int64(len(keyChars))))
+		if err != nil {
+			return nil, fmt.Errorf("generate jwt key: %w", err)
+		}
+		buf[i] = keyChars[n.Int64()]
+	}
+	if err := os.WriteFile(path, buf, 0600); err != nil {
+		return nil, fmt.Errorf("write jwt key: %w", err)
+	}
+	log.Printf("[auth] 新 JWT key 已写入 %s", path)
+	return buf, nil
+}
+
+// signJWT issues a 7-day HS256 token for user ARTEX.
+func signJWT(key []byte) (string, error) {
+	return jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.RegisteredClaims{
+		Subject:   "ARTEX",
+		ExpiresAt: jwt.NewNumericDate(time.Now().Add(jwtTTL)),
+		IssuedAt:  jwt.NewNumericDate(time.Now()),
+	}).SignedString(key)
+}
+
+// verifyJWT returns true when tokenStr is a valid, non-expired HS256 token.
+func verifyJWT(tokenStr string, key []byte) bool {
+	t, err := jwt.Parse(tokenStr, func(t *jwt.Token) (any, error) {
+		if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
+			return nil, fmt.Errorf("unexpected signing method")
+		}
+		return key, nil
+	})
+	return err == nil && t.Valid
+}
+
+// extractToken reads the JWT from Authorization: Bearer header,
+// artex_token cookie, or ?token= query param (for SSE connections).
+func extractToken(r *http.Request) string {
+	if h := r.Header.Get("Authorization"); strings.HasPrefix(h, "Bearer ") {
+		return strings.TrimPrefix(h, "Bearer ")
+	}
+	if c, err := r.Cookie("artex_token"); err == nil && c.Value != "" {
+		return c.Value
+	}
+	return r.URL.Query().Get("token")
+}
+
+// requireAuth wraps h with JWT validation.
+// /api/auth/* and /api/health are exempt.
+func (s *Server) requireAuth(h http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		p := r.URL.Path
+		if strings.HasPrefix(p, "/api/auth/") || p == "/api/health" {
+			h.ServeHTTP(w, r)
+			return
+		}
+		tok := extractToken(r)
+		if tok == "" {
+			writeErr(w, 401, "未授权")
+			return
+		}
+		if !verifyJWT(tok, s.jwtKey) {
+			writeErr(w, 401, "token 无效或已过期")
+			return
+		}
+		h.ServeHTTP(w, r)
+	})
+}
+
+// GET /api/auth/status — reports whether the admin password has been initialised.
+func (s *Server) authStatus(w http.ResponseWriter, r *http.Request) {
+	pg := s.pg(w)
+	if pg == nil {
+		return
+	}
+	hash, _, _ := pg.GetSetting(authPassKey)
+	writeJSON(w, 200, map[string]any{"initialized": hash != ""})
+}
+
+// POST /api/auth/init — sets the password for the first time; rejected if already set.
+func (s *Server) authInit(w http.ResponseWriter, r *http.Request) {
+	pg := s.pg(w)
+	if pg == nil {
+		return
+	}
+	existing, _, _ := pg.GetSetting(authPassKey)
+	if existing != "" {
+		writeErr(w, 403, "密码已设置")
+		return
+	}
+	var req struct {
+		Password string `json:"password"`
+	}
+	if err := decode(r, &req); err != nil || req.Password == "" {
+		writeErr(w, 400, "密码不能为空")
+		return
+	}
+	hash, err := bcrypt.GenerateFromPassword([]byte(req.Password), bcrypt.DefaultCost)
+	if err != nil {
+		writeErr(w, 500, "密码加密失败")
+		return
+	}
+	if err := pg.SetSetting(authPassKey, string(hash)); err != nil {
+		writeErr(w, 500, "保存失败: "+err.Error())
+		return
+	}
+	tok, err := signJWT(s.jwtKey)
+	if err != nil {
+		writeErr(w, 500, "token 生成失败")
+		return
+	}
+	writeJSON(w, 200, map[string]any{"token": tok})
+}
+
+// POST /api/auth/change-password — changes the admin password. Requires a valid
+// token (this route is under /api/auth/* which requireAuth exempts, so the token
+// is validated here) AND the current password.
+func (s *Server) authChangePassword(w http.ResponseWriter, r *http.Request) {
+	pg := s.pg(w)
+	if pg == nil {
+		return
+	}
+	if !verifyJWT(extractToken(r), s.jwtKey) {
+		writeErr(w, 401, "未授权")
+		return
+	}
+	var req struct {
+		OldPassword string `json:"old_password"`
+		NewPassword string `json:"new_password"`
+	}
+	if err := decode(r, &req); err != nil {
+		writeErr(w, 400, "请求格式错误")
+		return
+	}
+	if req.NewPassword == "" {
+		writeErr(w, 400, "新密码不能为空")
+		return
+	}
+	hash, ok, _ := pg.GetSetting(authPassKey)
+	if !ok || hash == "" {
+		writeErr(w, 403, "密码未初始化，请先设置密码")
+		return
+	}
+	if err := bcrypt.CompareHashAndPassword([]byte(hash), []byte(req.OldPassword)); err != nil {
+		writeErr(w, 401, "当前密码错误")
+		return
+	}
+	newHash, err := bcrypt.GenerateFromPassword([]byte(req.NewPassword), bcrypt.DefaultCost)
+	if err != nil {
+		writeErr(w, 500, "密码加密失败")
+		return
+	}
+	if err := pg.SetSetting(authPassKey, string(newHash)); err != nil {
+		writeErr(w, 500, "保存失败: "+err.Error())
+		return
+	}
+	writeJSON(w, 200, map[string]any{"ok": true})
+}
+
+// POST /api/auth/login — validates username/password and returns a JWT.
+func (s *Server) authLogin(w http.ResponseWriter, r *http.Request) {
+	pg := s.pg(w)
+	if pg == nil {
+		return
+	}
+	var req struct {
+		Username string `json:"username"`
+		Password string `json:"password"`
+	}
+	if err := decode(r, &req); err != nil {
+		writeErr(w, 400, "请求格式错误")
+		return
+	}
+	if req.Username != "ARTEX" {
+		writeErr(w, 401, "用户名或密码错误")
+		return
+	}
+	hash, ok, _ := pg.GetSetting(authPassKey)
+	if !ok || hash == "" {
+		writeErr(w, 403, "密码未初始化，请先设置密码")
+		return
+	}
+	if err := bcrypt.CompareHashAndPassword([]byte(hash), []byte(req.Password)); err != nil {
+		writeErr(w, 401, "用户名或密码错误")
+		return
+	}
+	tok, err := signJWT(s.jwtKey)
+	if err != nil {
+		writeErr(w, 500, "token 生成失败")
+		return
+	}
+	writeJSON(w, 200, map[string]any{"token": tok})
+}