Add files via upload

This commit is contained in:
公明
2026-07-11 11:44:25 +08:00
committed by GitHub
parent 142977413e
commit 3c6cf633e1
7 changed files with 143 additions and 360 deletions
+29 -70
View File
@@ -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 {
@@ -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)
}
}
+2 -5
View File
@@ -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")
+13 -13
View File
@@ -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)),
)
}
+24
View File
@@ -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
}