mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-07-14 07:57:26 +02:00
Add files via upload
This commit is contained in:
+9
-197
@@ -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 位十六进制)
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
|
||||
@@ -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)),
|
||||
)
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user