From 3c6cf633e1cf190599f94a3826cafa26b02ddfda Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Sat, 11 Jul 2026 11:44:25 +0800 Subject: [PATCH] Add files via upload --- internal/config/config.go | 206 +----------------- internal/config/config_test.go | 103 +++------ internal/security/auth_manager.go | 99 +++------ .../security/auth_manager_bootstrap_test.go | 38 ++++ internal/security/auth_manager_test.go | 7 +- internal/security/executor.go | 26 +-- internal/security/password.go | 24 ++ 7 files changed, 143 insertions(+), 360 deletions(-) create mode 100644 internal/security/auth_manager_bootstrap_test.go create mode 100644 internal/security/password.go diff --git a/internal/config/config.go b/internal/config/config.go index aed6ad75..06a253fa 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -2,7 +2,6 @@ package config import ( "crypto/rand" - "encoding/base64" "encoding/hex" "encoding/json" "fmt" @@ -12,6 +11,8 @@ import ( "strconv" "strings" + "cyberstrike-ai/internal/termout" + "gopkg.in/yaml.v3" ) @@ -43,9 +44,8 @@ type Config struct { } type EnsureLocalConfigResult struct { - Created bool - GeneratedPassword string - ExamplePath string + Created bool + ExamplePath string } const ( @@ -1005,11 +1005,7 @@ func normalizeHitlModeForPrompt(mode string) string { } type AuthConfig struct { - Password string `yaml:"password" json:"password"` - SessionDurationHours int `yaml:"session_duration_hours" json:"session_duration_hours"` - GeneratedPassword string `yaml:"-" json:"-"` - GeneratedPasswordPersisted bool `yaml:"-" json:"-"` - GeneratedPasswordPersistErr string `yaml:"-" json:"-"` + SessionDurationHours int `yaml:"session_duration_hours" json:"session_duration_hours"` } // MonitorConfig MCP 状态监控（tool_executions）保留策略。 @@ -1169,23 +1165,6 @@ func Load(path string) (*Config, error) { if cfg.Audit.MaxDetailBytes <= 0 { cfg.Audit.MaxDetailBytes = 8192 } - if strings.TrimSpace(cfg.Auth.Password) == "" { - password, err := generateStrongPassword(24) - if err != nil { - return nil, fmt.Errorf("生成默认密码失败: %w", err) - } - - cfg.Auth.Password = password - cfg.Auth.GeneratedPassword = password - - if err := PersistAuthPassword(path, password); err != nil { - cfg.Auth.GeneratedPasswordPersisted = false - cfg.Auth.GeneratedPasswordPersistErr = err.Error() - } else { - cfg.Auth.GeneratedPasswordPersisted = true - } - } - // 如果配置了工具目录，从目录加载工具配置 if cfg.Security.ToolsDir != "" { inlineTools := append([]ToolConfig(nil), cfg.Security.Tools...) @@ -1291,181 +1270,14 @@ func EnsureLocalConfig(path string) (EnsureLocalConfigResult, error) { return EnsureLocalConfigResult{}, fmt.Errorf("创建配置文件失败: %w", err) } - password, err := generateStrongPassword(24) - if err != nil { - return EnsureLocalConfigResult{}, fmt.Errorf("生成默认密码失败: %w", err) - } - if err := PersistAuthPassword(path, password); err != nil { - return EnsureLocalConfigResult{}, fmt.Errorf("写入默认密码失败: %w", err) - } - return EnsureLocalConfigResult{ - Created: true, - GeneratedPassword: password, - ExamplePath: examplePath, + Created: true, + ExamplePath: examplePath, }, nil } -func generateStrongPassword(length int) (string, error) { - if length <= 0 { - length = 24 - } - - bytesLen := length - randomBytes := make([]byte, bytesLen) - if _, err := rand.Read(randomBytes); err != nil { - return "", err - } - - password := base64.RawURLEncoding.EncodeToString(randomBytes) - if len(password) > length { - password = password[:length] - } - return password, nil -} - -func PersistAuthPassword(path, password string) error { - data, err := os.ReadFile(path) - if err != nil { - return err - } - - lines := strings.Split(string(data), "\n") - inAuthBlock := false - authIndent := -1 - - for i, line := range lines { - trimmed := strings.TrimSpace(line) - if !inAuthBlock { - if strings.HasPrefix(trimmed, "auth:") { - inAuthBlock = true - authIndent = len(line) - len(strings.TrimLeft(line, " ")) - } - continue - } - - if trimmed == "" || strings.HasPrefix(trimmed, "#") { - continue - } - - leadingSpaces := len(line) - len(strings.TrimLeft(line, " ")) - if leadingSpaces <= authIndent { - // 离开 auth 块 - inAuthBlock = false - authIndent = -1 - // 继续寻找其它 auth 块（理论上没有） - if strings.HasPrefix(trimmed, "auth:") { - inAuthBlock = true - authIndent = leadingSpaces - } - continue - } - - if strings.HasPrefix(strings.TrimSpace(line), "password:") { - prefix := line[:len(line)-len(strings.TrimLeft(line, " "))] - comment := "" - if idx := yamlLineCommentIndex(line); idx >= 0 { - comment = strings.TrimRight(line[idx:], " ") - } - - newLine := fmt.Sprintf("%spassword: %s", prefix, quoteYAMLString(password)) - if comment != "" { - if !strings.HasPrefix(comment, " ") { - newLine += " " - } - newLine += comment - } - lines[i] = newLine - break - } - } - - return os.WriteFile(path, []byte(strings.Join(lines, "\n")), 0644) -} - -func quoteYAMLString(value string) string { - node := yaml.Node{ - Kind: yaml.ScalarNode, - Tag: "!!str", - Style: yaml.DoubleQuotedStyle, - Value: value, - } - data, err := yaml.Marshal(&node) - if err != nil { - return strconv.Quote(value) - } - return strings.TrimSuffix(string(data), "\n") -} - -func yamlLineCommentIndex(line string) int { - inSingleQuote := false - inDoubleQuote := false - escaped := false - - for i, r := range line { - if inDoubleQuote { - if escaped { - escaped = false - continue - } - if r == '\\' { - escaped = true - continue - } - if r == '"' { - inDoubleQuote = false - } - continue - } - if inSingleQuote { - if r == '\'' { - inSingleQuote = false - } - continue - } - - switch r { - case '"': - inDoubleQuote = true - case '\'': - inSingleQuote = true - case '#': - if i == 0 || isYAMLWhitespace(line[i-1]) { - return i - } - } - } - return -1 -} - -func isYAMLWhitespace(b byte) bool { - return b == ' ' || b == '\t' -} - -func PrintGeneratedPasswordWarning(password string, persisted bool, persistErr string) { - if strings.TrimSpace(password) == "" { - return - } - - if persisted { - fmt.Println("[CyberStrikeAI] ✅ 已为您自动生成并写入 Web 登录密码。") - } else { - if persistErr != "" { - fmt.Printf("[CyberStrikeAI] ⚠️ 无法自动写入配置文件中的密码: %s\n", persistErr) - } else { - fmt.Println("[CyberStrikeAI] ⚠️ 无法自动写入配置文件中的密码。") - } - fmt.Println("请手动将以下随机密码写入 config.yaml 的 auth.password：") - } - - fmt.Println("----------------------------------------------------------------") - fmt.Println("CyberStrikeAI Auto-Generated Web Password") - fmt.Printf("Password: %s\n", password) - fmt.Println("WARNING: Anyone with this password can fully control CyberStrikeAI.") - fmt.Println("Please store it securely and change it in config.yaml as soon as possible.") - fmt.Println("警告：持有此密码的人将拥有对 CyberStrikeAI 的完全控制权限。") - fmt.Println("请妥善保管，并尽快在 config.yaml 中修改 auth.password！") - fmt.Println("----------------------------------------------------------------") +func PrintBootstrapAdminPassword(password string) { + termout.PrintBootstrapAdminCredentials(password) } // generateRandomToken 生成用于 MCP 鉴权的随机字符串（64 位十六进制） diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 772852ff..bbbadf69 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -7,78 +7,12 @@ import ( "testing" ) -func TestPersistAuthPasswordQuotesYAMLSpecialCharacters(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "config.yaml") - initial := strings.Join([]string{ - "server:", - " host: 0.0.0.0", - "auth:", - " password: old-password # Web 登录密码", - " session_duration_hours: 12", - "log:", - " level: info", - "", - }, "\n") - if err := os.WriteFile(path, []byte(initial), 0644); err != nil { - t.Fatalf("write config: %v", err) - } - - want := `@abc:def # still password` - if err := PersistAuthPassword(path, want); err != nil { - t.Fatalf("PersistAuthPassword: %v", err) - } - - data, err := os.ReadFile(path) - if err != nil { - t.Fatalf("read config: %v", err) - } - if !strings.Contains(string(data), `password: "@abc:def # still password" # Web 登录密码`) { - t.Fatalf("password was not safely quoted or comment was not preserved:\n%s", data) - } - - cfg, err := Load(path) - if err != nil { - t.Fatalf("Load after PersistAuthPassword: %v", err) - } - if cfg.Auth.Password != want { - t.Fatalf("Auth.Password = %q, want %q", cfg.Auth.Password, want) - } -} - -func TestPersistAuthPasswordDoesNotTreatQuotedHashAsComment(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "config.yaml") - initial := strings.Join([]string{ - "auth:", - ` password: "old#password"`, - " session_duration_hours: 12", - "", - }, "\n") - if err := os.WriteFile(path, []byte(initial), 0644); err != nil { - t.Fatalf("write config: %v", err) - } - - if err := PersistAuthPassword(path, "new-password"); err != nil { - t.Fatalf("PersistAuthPassword: %v", err) - } - - data, err := os.ReadFile(path) - if err != nil { - t.Fatalf("read config: %v", err) - } - if strings.Contains(string(data), "#password") { - t.Fatalf("old quoted password fragment was incorrectly preserved as a comment:\n%s", data) - } -} - -func TestEnsureLocalConfigCreatesFromExampleWithGeneratedPassword(t *testing.T) { +func TestEnsureLocalConfigCreatesFromExample(t *testing.T) { dir := t.TempDir() examplePath := filepath.Join(dir, "config.example.yaml") configPath := filepath.Join(dir, "config.yaml") example := []byte(`auth: - password: "change-me-use-a-long-random-password" session_duration_hours: 12 server: host: 127.0.0.1 @@ -95,9 +29,6 @@ server: if !result.Created { t.Fatal("Created = false, want true") } - if result.GeneratedPassword == "" { - t.Fatal("GeneratedPassword is empty") - } if result.ExamplePath != examplePath { t.Fatalf("ExamplePath = %q, want %q", result.ExamplePath, examplePath) } @@ -106,11 +37,8 @@ server: if err != nil { t.Fatalf("Load generated config: %v", err) } - if cfg.Auth.Password == "change-me-use-a-long-random-password" { - t.Fatal("auth.password still contains the template placeholder") - } - if cfg.Auth.Password != result.GeneratedPassword { - t.Fatalf("Auth.Password = %q, want generated password %q", cfg.Auth.Password, result.GeneratedPassword) + if cfg.Auth.SessionDurationHours != 12 { + t.Fatalf("SessionDurationHours = %d, want 12", cfg.Auth.SessionDurationHours) } second, err := EnsureLocalConfig(configPath) @@ -122,6 +50,31 @@ server: } } +func TestLoadIgnoresLegacyAuthPasswordField(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "config.yaml") + initial := strings.Join([]string{ + "auth:", + ` password: "legacy-password"`, + " session_duration_hours: 12", + "server:", + " host: 127.0.0.1", + " port: 8080", + "", + }, "\n") + if err := os.WriteFile(path, []byte(initial), 0644); err != nil { + t.Fatalf("write config: %v", err) + } + + cfg, err := Load(path) + if err != nil { + t.Fatalf("Load: %v", err) + } + if cfg.Auth.SessionDurationHours != 12 { + t.Fatalf("SessionDurationHours = %d, want 12", cfg.Auth.SessionDurationHours) + } +} + func TestHitlAuditModelEffectiveFallsBackToMainConfig(t *testing.T) { main := OpenAIConfig{ Provider: "openai", diff --git a/internal/security/auth_manager.go b/internal/security/auth_manager.go index 7ec2d5b1..74e3b79a 100644 --- a/internal/security/auth_manager.go +++ b/internal/security/auth_manager.go @@ -32,7 +32,6 @@ type Session struct { // AuthManager manages password-based authentication and session lifecycle. type AuthManager struct { - password string sessionDuration time.Duration db *database.DB @@ -41,39 +40,49 @@ type AuthManager struct { } // NewAuthManager creates a new AuthManager instance. -func NewAuthManager(password string, sessionDurationHours int) (*AuthManager, error) { - if strings.TrimSpace(password) == "" { - return nil, errors.New("auth password must be configured") - } - +func NewAuthManager(sessionDurationHours int) *AuthManager { if sessionDurationHours <= 0 { sessionDurationHours = 12 } return &AuthManager{ - password: password, sessionDuration: time.Duration(sessionDurationHours) * time.Hour, sessions: make(map[string]Session), - }, nil + } } -// AttachRBACStore enables multi-user RBAC authentication and bootstraps the -// built-in admin account from the legacy auth.password value. -func (a *AuthManager) AttachRBACStore(db *database.DB) error { +// AttachRBACStore enables multi-user RBAC authentication. When no users exist yet, +// it bootstraps the built-in admin account and returns the generated initial password. +func (a *AuthManager) AttachRBACStore(db *database.DB) (generatedAdminPassword string, err error) { if db == nil { - return nil + return "", errors.New("database is required for authentication") } - hash, err := HashPassword(a.password) + + needsAdminPassword, err := db.RBACNeedsAdminPassword() if err != nil { - return err + return "", err } - if err := db.BootstrapRBAC(hash, PermissionCatalog); err != nil { - return err + + adminPasswordHash := "" + if needsAdminPassword { + generatedAdminPassword, err = GenerateStrongPassword(24) + if err != nil { + return "", err + } + adminPasswordHash, err = HashPassword(generatedAdminPassword) + if err != nil { + return "", err + } } + + if err := db.BootstrapRBAC(adminPasswordHash, PermissionCatalog); err != nil { + return "", err + } + a.mu.Lock() a.db = db a.mu.Unlock() - return nil + return generatedAdminPassword, nil } // Authenticate validates the password and creates a new session. @@ -94,23 +103,9 @@ func (a *AuthManager) authenticateSession(username, password string) (Session, e a.mu.RLock() db := a.db - legacyPassword := a.password a.mu.RUnlock() - if db == nil { - if password != legacyPassword { - return Session{}, ErrInvalidPassword - } - return Session{ - Token: token, - ExpiresAt: expiresAt, - UserID: "admin", - Username: "admin", - DisplayName: "管理员", - Roles: []string{database.RBACSystemRoleAdmin}, - Permissions: allPermissions(), - Scope: database.RBACScopeAll, - }, nil + return Session{}, errors.New("authentication store is not configured") } username = strings.TrimSpace(strings.ToLower(username)) @@ -187,10 +182,9 @@ func (a *AuthManager) CheckPassword(password string) bool { func (a *AuthManager) CheckUserPassword(username, password string) bool { a.mu.RLock() db := a.db - legacyPassword := a.password a.mu.RUnlock() if db == nil { - return password == legacyPassword + return false } user, err := db.GetRBACUserByUsername(username) if err != nil { @@ -212,7 +206,7 @@ func (a *AuthManager) UpdateUserPassword(userID, password string) error { db := a.db a.mu.RUnlock() if db == nil { - return a.UpdateConfig(password, a.SessionDurationHours()) + return errors.New("authentication store is not configured") } if err := db.UpdateRBACUserPassword(userID, hash); err != nil { return err @@ -263,41 +257,6 @@ func (a *AuthManager) SessionDurationHours() int { return int(a.sessionDuration / time.Hour) } -// UpdateConfig updates the password and session duration, revoking existing sessions. -func (a *AuthManager) UpdateConfig(password string, sessionDurationHours int) error { - password = strings.TrimSpace(password) - if password == "" { - return errors.New("auth password must be configured") - } - - if sessionDurationHours <= 0 { - sessionDurationHours = 12 - } - - hash := "" - if a.db != nil { - var err error - hash, err = HashPassword(password) - if err != nil { - return err - } - } - - a.mu.Lock() - a.password = password - a.sessionDuration = time.Duration(sessionDurationHours) * time.Hour - a.sessions = make(map[string]Session) - db := a.db - a.mu.Unlock() - - if db != nil { - if err := db.UpdateRBACAdminPassword(hash); err != nil { - return err - } - } - return nil -} - func allPermissions() map[string]bool { out := make(map[string]bool, len(PermissionCatalog)) for key := range PermissionCatalog { diff --git a/internal/security/auth_manager_bootstrap_test.go b/internal/security/auth_manager_bootstrap_test.go new file mode 100644 index 00000000..fff311b3 --- /dev/null +++ b/internal/security/auth_manager_bootstrap_test.go @@ -0,0 +1,38 @@ +package security + +import ( + "path/filepath" + "testing" + + "cyberstrike-ai/internal/database" + + "go.uber.org/zap" +) + +func TestAttachRBACStoreBootstrapsAdminPassword(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "auth-bootstrap.db"), zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + + manager := NewAuthManager(12) + generated, err := manager.AttachRBACStore(db) + if err != nil { + t.Fatalf("AttachRBACStore: %v", err) + } + if generated == "" { + t.Fatal("expected generated admin password on first bootstrap") + } + if !manager.CheckUserPassword("admin", generated) { + t.Fatal("generated password should authenticate admin") + } + + second, err := manager.AttachRBACStore(db) + if err != nil { + t.Fatalf("AttachRBACStore second call: %v", err) + } + if second != "" { + t.Fatalf("expected no password on second bootstrap, got %q", second) + } +} diff --git a/internal/security/auth_manager_test.go b/internal/security/auth_manager_test.go index 83456cc8..25e1b591 100644 --- a/internal/security/auth_manager_test.go +++ b/internal/security/auth_manager_test.go @@ -20,11 +20,8 @@ func TestAuthManagerAuthenticatesCreatedRBACUser(t *testing.T) { } t.Cleanup(func() { _ = db.Close() }) - manager, err := NewAuthManager("admin-secret", 12) - if err != nil { - t.Fatalf("NewAuthManager: %v", err) - } - if err := manager.AttachRBACStore(db); err != nil { + manager := NewAuthManager(12) + if _, err := manager.AttachRBACStore(db); err != nil { t.Fatalf("AttachRBACStore: %v", err) } hash, err := HashPassword("operator-secret") diff --git a/internal/security/executor.go b/internal/security/executor.go index 35b95b3a..ce402cb8 100644 --- a/internal/security/executor.go +++ b/internal/security/executor.go @@ -65,7 +65,7 @@ func (e *Executor) buildToolIndex() { e.toolIndex[e.config.Tools[i].Name] = &e.config.Tools[i] } } - e.logger.Info("工具索引构建完成", + e.logger.Debug("工具索引构建完成", zap.Int("totalTools", len(e.config.Tools)), zap.Int("enabledTools", len(e.toolIndex)), ) @@ -73,14 +73,14 @@ func (e *Executor) buildToolIndex() { // ExecuteTool 执行安全工具 func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[string]interface{}) (*mcp.ToolResult, error) { - e.logger.Info("ExecuteTool被调用", + e.logger.Debug("ExecuteTool被调用", zap.String("toolName", toolName), zap.Any("args", args), ) // 特殊处理：exec工具直接执行系统命令 if toolName == "exec" { - e.logger.Info("执行exec工具") + e.logger.Debug("执行exec工具") return e.executeSystemCommand(ctx, args) } @@ -95,7 +95,7 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st return nil, fmt.Errorf("工具 %s 未找到或未启用", toolName) } - e.logger.Info("找到工具配置", + e.logger.Debug("找到工具配置", zap.String("toolName", toolName), zap.String("command", toolConfig.Command), zap.Strings("args", toolConfig.Args), @@ -103,7 +103,7 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st // 特殊处理：内部工具（command 以 "internal:" 开头） if strings.HasPrefix(toolConfig.Command, "internal:") { - e.logger.Info("执行内部工具", + e.logger.Debug("执行内部工具", zap.String("toolName", toolName), zap.String("command", toolConfig.Command), ) @@ -113,7 +113,7 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st // 构建命令 - 根据工具类型使用不同的参数格式 cmdArgs := e.buildCommandArgs(toolName, toolConfig, args) - e.logger.Info("构建命令参数完成", + e.logger.Debug("构建命令参数完成", zap.String("toolName", toolName), zap.Strings("cmdArgs", cmdArgs), zap.Int("argsCount", len(cmdArgs)), @@ -142,7 +142,7 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st attachNonInteractiveStdin(cmd) _ = prepareShellCmdSession(cmd) - e.logger.Info("执行安全工具", + e.logger.Debug("执行安全工具", zap.String("tool", toolName), zap.Strings("args", cmdArgs), ) @@ -180,7 +180,7 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st if exitCode != nil && toolConfig.AllowedExitCodes != nil { for _, allowedCode := range toolConfig.AllowedExitCodes { if *exitCode == allowedCode { - e.logger.Info("工具执行完成（退出码在允许列表中）", + e.logger.Debug("工具执行完成（退出码在允许列表中）", zap.String("tool", toolName), zap.Int("exitCode", *exitCode), zap.String("output", string(output)), @@ -215,7 +215,7 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st }, nil } - e.logger.Info("工具执行成功", + e.logger.Debug("工具执行成功", zap.String("tool", toolName), zap.String("output", string(output)), ) @@ -233,7 +233,7 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st // RegisterTools 注册工具到MCP服务器 func (e *Executor) RegisterTools(mcpServer *mcp.Server) { - e.logger.Info("开始注册工具", + e.logger.Debug("开始注册工具", zap.Int("totalTools", len(e.config.Tools)), zap.Int("enabledTools", len(e.toolIndex)), ) @@ -281,7 +281,7 @@ func (e *Executor) RegisterTools(mcpServer *mcp.Server) { } handler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { - e.logger.Info("工具handler被调用", + e.logger.Debug("工具handler被调用", zap.String("toolName", toolName), zap.Any("args", args), ) @@ -289,14 +289,14 @@ func (e *Executor) RegisterTools(mcpServer *mcp.Server) { } mcpServer.RegisterTool(tool, handler) - e.logger.Info("注册安全工具成功", + e.logger.Debug("注册安全工具成功", zap.String("tool", toolConfigCopy.Name), zap.String("command", toolConfigCopy.Command), zap.Int("index", i), ) } - e.logger.Info("工具注册完成", + e.logger.Debug("工具注册完成", zap.Int("registeredCount", len(e.config.Tools)), ) } diff --git a/internal/security/password.go b/internal/security/password.go new file mode 100644 index 00000000..3eb356de --- /dev/null +++ b/internal/security/password.go @@ -0,0 +1,24 @@ +package security + +import ( + "crypto/rand" + "encoding/base64" +) + +// GenerateStrongPassword returns a URL-safe random password of the given length. +func GenerateStrongPassword(length int) (string, error) { + if length <= 0 { + length = 24 + } + + randomBytes := make([]byte, length) + if _, err := rand.Read(randomBytes); err != nil { + return "", err + } + + password := base64.RawURLEncoding.EncodeToString(randomBytes) + if len(password) > length { + password = password[:length] + } + return password, nil +}