diff --git a/internal/app/app.go b/internal/app/app.go index 213edce1..070e0e53 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -109,6 +109,12 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error if err := authManager.AttachRBACStore(db); err != nil { return nil, fmt.Errorf("初始化RBAC失败: %w", err) } + for platform, userID := range cfg.Robots.ServiceAccountUserIDs() { + user, userErr := db.GetRBACUserByID(userID) + if userErr != nil || !user.Enabled { + return nil, fmt.Errorf("robots.%s.auth.service_user_id 必须指向已启用的 RBAC 用户", platform) + } + } auditSvc := audit.NewService(db, cfg, log.Logger) audit.RegisterConversationCreateHook(auditSvc) @@ -381,6 +387,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error registerWebshellTools(mcpServer, db, webshellHandler, log.Logger) registerWebshellManagementTools(mcpServer, db, webshellHandler, log.Logger) configHandler := handler.NewConfigHandler(configPath, cfg, mcpServer, executor, agent, attackChainHandler, externalMCPMgr, log.Logger) + configHandler.SetDB(db) configHandler.SetAudit(auditSvc) agentHandler.SetHitlToolWhitelistSaver(configHandler) agentHandler.SetHitlAuditStrategySaver(configHandler) @@ -413,6 +420,7 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error conversationHandler.SetTaskStopper(agentHandler) auditHandler := handler.NewAuditHandler(db, auditSvc, log.Logger) robotHandler := handler.NewRobotHandler(cfg, db, agentHandler, log.Logger) + robotHandler.SetAudit(auditSvc) openAPIHandler := handler.NewOpenAPIHandler(db, log.Logger, conversationHandler, agentHandler) // 创建 App 实例(部分字段稍后填充) @@ -860,6 +868,9 @@ func setupRoutes( authRoutes.POST("/logout", security.AuthMiddleware(authManager), authHandler.Logout) authRoutes.POST("/change-password", security.AuthMiddleware(authManager), security.RequirePermission("auth:self"), authHandler.ChangePassword) authRoutes.GET("/validate", security.AuthMiddleware(authManager), authHandler.Validate) + authRoutes.POST("/robot-binding-code", security.AuthMiddleware(authManager), security.RequirePermission("auth:self"), robotHandler.CreateRobotBindingCode) + authRoutes.GET("/robot-bindings", security.AuthMiddleware(authManager), security.RequirePermission("auth:self"), robotHandler.ListMyRobotBindings) + authRoutes.DELETE("/robot-bindings/:id", security.AuthMiddleware(authManager), security.RequirePermission("auth:self"), robotHandler.DeleteMyRobotBinding) } // 机器人回调(无需登录,供企业微信/钉钉/飞书服务器调用) diff --git a/internal/database/rbac.go b/internal/database/rbac.go index de2ed328..53b1716b 100644 --- a/internal/database/rbac.go +++ b/internal/database/rbac.go @@ -137,6 +137,25 @@ func (db *DB) initRBACTables() error { FOREIGN KEY (user_id) REFERENCES rbac_users(id) ON DELETE CASCADE, UNIQUE(user_id, resource_type, resource_id) );`, + `CREATE TABLE IF NOT EXISTS robot_user_bindings ( + id TEXT PRIMARY KEY, + platform TEXT NOT NULL, + external_user_id TEXT NOT NULL, + rbac_user_id TEXT NOT NULL, + enabled INTEGER NOT NULL DEFAULT 1, + created_at DATETIME NOT NULL, + updated_at DATETIME NOT NULL, + FOREIGN KEY (rbac_user_id) REFERENCES rbac_users(id) ON DELETE CASCADE, + UNIQUE(platform, external_user_id) + );`, + `CREATE TABLE IF NOT EXISTS robot_binding_codes ( + code_hash TEXT PRIMARY KEY, + rbac_user_id TEXT NOT NULL, + expires_at DATETIME NOT NULL, + used_at DATETIME, + created_at DATETIME NOT NULL, + FOREIGN KEY (rbac_user_id) REFERENCES rbac_users(id) ON DELETE CASCADE + );`, `CREATE TABLE IF NOT EXISTS chat_upload_artifacts ( relative_path TEXT PRIMARY KEY, conversation_id TEXT NOT NULL, @@ -155,6 +174,8 @@ func (db *DB) initRBACTables() error { `CREATE INDEX IF NOT EXISTS idx_rbac_role_permissions_role ON rbac_role_permissions(role_id);`, `CREATE INDEX IF NOT EXISTS idx_rbac_assignments_user_resource ON rbac_resource_assignments(user_id, resource_type, resource_id);`, `CREATE INDEX IF NOT EXISTS idx_rbac_assignments_resource ON rbac_resource_assignments(resource_type, resource_id);`, + `CREATE INDEX IF NOT EXISTS idx_robot_user_bindings_user ON robot_user_bindings(rbac_user_id);`, + `CREATE INDEX IF NOT EXISTS idx_robot_binding_codes_expiry ON robot_binding_codes(expires_at);`, `CREATE INDEX IF NOT EXISTS idx_chat_upload_artifacts_conversation ON chat_upload_artifacts(conversation_id);`, `CREATE INDEX IF NOT EXISTS idx_chat_upload_artifacts_owner ON chat_upload_artifacts(owner_user_id);`, `CREATE INDEX IF NOT EXISTS idx_c2_payload_artifacts_listener ON c2_payload_artifacts(listener_id);`, diff --git a/internal/database/robot_identity.go b/internal/database/robot_identity.go new file mode 100644 index 00000000..a7e0039e --- /dev/null +++ b/internal/database/robot_identity.go @@ -0,0 +1,174 @@ +package database + +import ( + "database/sql" + "fmt" + "strings" + "time" + + "github.com/google/uuid" +) + +// RobotUserBinding maps one tenant-scoped platform identity to one RBAC user. +// external_user_id must be derived from the verified platform event, never +// from user-controlled message content. +type RobotUserBinding struct { + ID string `json:"id"` + Platform string `json:"platform"` + ExternalUserID string `json:"externalUserId"` + RBACUserID string `json:"rbacUserId"` + Enabled bool `json:"enabled"` + CreatedAt time.Time `json:"createdAt"` + UpdatedAt time.Time `json:"updatedAt"` +} + +func normalizeRobotIdentity(platform, externalUserID string) (string, string, error) { + platform = strings.ToLower(strings.TrimSpace(platform)) + externalUserID = strings.TrimSpace(externalUserID) + if platform == "" || externalUserID == "" { + return "", "", fmt.Errorf("robot platform and external user identity are required") + } + return platform, externalUserID, nil +} + +func (db *DB) CreateRobotBindingCode(userID, codeHash string, expiresAt time.Time) error { + userID = strings.TrimSpace(userID) + codeHash = strings.TrimSpace(codeHash) + if userID == "" || codeHash == "" || !expiresAt.After(time.Now()) { + return fmt.Errorf("invalid robot binding code") + } + now := time.Now() + tx, err := db.Begin() + if err != nil { + return err + } + defer func() { _ = tx.Rollback() }() + // Keep only the newest active code per user and remove expired/used secrets. + if _, err = tx.Exec(`DELETE FROM robot_binding_codes WHERE rbac_user_id = ? OR expires_at <= ? OR used_at IS NOT NULL`, userID, now); err != nil { + return err + } + if _, err = tx.Exec(`INSERT INTO robot_binding_codes (code_hash, rbac_user_id, expires_at, created_at) VALUES (?, ?, ?, ?)`, codeHash, userID, expiresAt, now); err != nil { + return err + } + return tx.Commit() +} + +// ConsumeRobotBindingCode atomically consumes a single-use code and binds the +// verified platform identity. Existing bindings are deliberately replaced so +// users can recover from stale or incorrect associations with a fresh code. +func (db *DB) ConsumeRobotBindingCode(platform, externalUserID, codeHash string) (*RBACUser, error) { + platform, externalUserID, err := normalizeRobotIdentity(platform, externalUserID) + if err != nil { + return nil, err + } + codeHash = strings.TrimSpace(codeHash) + if codeHash == "" { + return nil, fmt.Errorf("binding code is required") + } + tx, err := db.Begin() + if err != nil { + return nil, err + } + defer func() { _ = tx.Rollback() }() + + var userID string + now := time.Now() + if err = tx.QueryRow(` + SELECT c.rbac_user_id + FROM robot_binding_codes c + JOIN rbac_users u ON u.id = c.rbac_user_id + WHERE c.code_hash = ? AND c.used_at IS NULL AND c.expires_at > ? AND u.enabled = 1 + `, codeHash, now).Scan(&userID); err != nil { + if err == sql.ErrNoRows { + return nil, fmt.Errorf("binding code is invalid or expired") + } + return nil, err + } + result, err := tx.Exec(`UPDATE robot_binding_codes SET used_at = ? WHERE code_hash = ? AND used_at IS NULL`, now, codeHash) + if err != nil { + return nil, err + } + if affected, _ := result.RowsAffected(); affected != 1 { + return nil, fmt.Errorf("binding code has already been used") + } + if _, err = tx.Exec(` + INSERT INTO robot_user_bindings (id, platform, external_user_id, rbac_user_id, enabled, created_at, updated_at) + VALUES (?, ?, ?, ?, 1, ?, ?) + ON CONFLICT(platform, external_user_id) DO UPDATE SET + rbac_user_id = excluded.rbac_user_id, + enabled = 1, + updated_at = excluded.updated_at + `, uuid.New().String(), platform, externalUserID, userID, now, now); err != nil { + return nil, err + } + if err = tx.Commit(); err != nil { + return nil, err + } + return db.GetRBACUserByID(userID) +} + +func (db *DB) ResolveRobotRBACAccess(platform, externalUserID string) (*RBACAccess, error) { + platform, externalUserID, err := normalizeRobotIdentity(platform, externalUserID) + if err != nil { + return nil, err + } + var userID string + err = db.QueryRow(` + SELECT b.rbac_user_id + FROM robot_user_bindings b + JOIN rbac_users u ON u.id = b.rbac_user_id + WHERE b.platform = ? AND b.external_user_id = ? AND b.enabled = 1 AND u.enabled = 1 + `, platform, externalUserID).Scan(&userID) + if err == sql.ErrNoRows { + return nil, fmt.Errorf("robot identity is not bound") + } + if err != nil { + return nil, err + } + return db.ResolveRBACAccess(userID) +} + +func (db *DB) ListRobotUserBindings(userID string) ([]RobotUserBinding, error) { + rows, err := db.Query(` + SELECT id, platform, external_user_id, rbac_user_id, enabled, created_at, updated_at + FROM robot_user_bindings WHERE rbac_user_id = ? ORDER BY updated_at DESC + `, strings.TrimSpace(userID)) + if err != nil { + return nil, err + } + defer rows.Close() + var out []RobotUserBinding + for rows.Next() { + var b RobotUserBinding + var enabled int + var createdAt, updatedAt string + if err := rows.Scan(&b.ID, &b.Platform, &b.ExternalUserID, &b.RBACUserID, &enabled, &createdAt, &updatedAt); err != nil { + return nil, err + } + b.Enabled = enabled != 0 + b.CreatedAt = parseDBTime(createdAt) + b.UpdatedAt = parseDBTime(updatedAt) + out = append(out, b) + } + return out, rows.Err() +} + +func (db *DB) DeleteRobotUserBindingForUser(bindingID, userID string) error { + result, err := db.Exec(`DELETE FROM robot_user_bindings WHERE id = ? AND rbac_user_id = ?`, strings.TrimSpace(bindingID), strings.TrimSpace(userID)) + if err != nil { + return err + } + if affected, _ := result.RowsAffected(); affected != 1 { + return sql.ErrNoRows + } + return nil +} + +func (db *DB) DeleteRobotIdentityBinding(platform, externalUserID string) error { + platform, externalUserID, err := normalizeRobotIdentity(platform, externalUserID) + if err != nil { + return err + } + _, err = db.Exec(`DELETE FROM robot_user_bindings WHERE platform = ? AND external_user_id = ?`, platform, externalUserID) + return err +} diff --git a/internal/database/robot_identity_test.go b/internal/database/robot_identity_test.go new file mode 100644 index 00000000..60a70c46 --- /dev/null +++ b/internal/database/robot_identity_test.go @@ -0,0 +1,89 @@ +package database_test + +import ( + "testing" + "time" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "go.uber.org/zap" +) + +func TestRobotBindingCodeIsSingleUseAndPermissionsAreResolvedLive(t *testing.T) { + db, err := database.NewDB(t.TempDir()+"/robot-identity.db", zap.NewNop()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + if err := db.BootstrapRBAC("hash", security.PermissionCatalog); err != nil { + t.Fatal(err) + } + user, err := db.CreateRBACUser("bound-user", "Bound User", "hash", true, []string{database.RBACSystemRoleOperator}) + if err != nil { + t.Fatal(err) + } + if err := db.CreateRobotBindingCode(user.ID, "code-hash", time.Now().Add(time.Minute)); err != nil { + t.Fatal(err) + } + bound, err := db.ConsumeRobotBindingCode("LARK", "t:tenant|u:user", "code-hash") + if err != nil || bound.ID != user.ID { + t.Fatalf("consume binding code: user=%v err=%v", bound, err) + } + if _, err := db.ConsumeRobotBindingCode("lark", "t:tenant|u:other", "code-hash"); err == nil { + t.Fatal("single-use binding code was accepted twice") + } + access, err := db.ResolveRobotRBACAccess("lark", "t:tenant|u:user") + if err != nil || !access.Permissions["agent:execute"] { + t.Fatalf("resolved access does not include live role permissions: %#v err=%v", access, err) + } + disabled := false + if err := db.UpdateRBACUser(user.ID, user.DisplayName, &disabled, nil); err != nil { + t.Fatal(err) + } + if _, err := db.ResolveRobotRBACAccess("lark", "t:tenant|u:user"); err == nil { + t.Fatal("disabled RBAC user retained robot access") + } +} + +func TestRobotBindingCodeExpiryAndOwnerScopedRevocation(t *testing.T) { + db, err := database.NewDB(t.TempDir()+"/robot-revoke.db", zap.NewNop()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + if err := db.BootstrapRBAC("hash", security.PermissionCatalog); err != nil { + t.Fatal(err) + } + u1, _ := db.CreateRBACUser("binding-owner", "Owner", "hash", true, nil) + u2, _ := db.CreateRBACUser("binding-other", "Other", "hash", true, nil) + now := time.Now() + if _, err := db.Exec(`INSERT INTO robot_binding_codes (code_hash, rbac_user_id, expires_at, created_at) VALUES (?, ?, ?, ?)`, "expired-hash", u1.ID, now.Add(-time.Minute), now.Add(-2*time.Minute)); err != nil { + t.Fatal(err) + } + if _, err := db.ConsumeRobotBindingCode("wecom", "t:corp|u:expired", "expired-hash"); err == nil { + t.Fatal("expired binding code was accepted") + } + if err := db.CreateRobotBindingCode(u1.ID, "valid-hash", time.Now().Add(time.Minute)); err != nil { + t.Fatal(err) + } + if _, err := db.ConsumeRobotBindingCode("wecom", "t:corp|u:one", "valid-hash"); err != nil { + t.Fatal(err) + } + bindings, err := db.ListRobotUserBindings(u1.ID) + if err != nil || len(bindings) != 1 { + t.Fatalf("bindings=%v err=%v", bindings, err) + } + if err := db.DeleteRobotUserBindingForUser(bindings[0].ID, u2.ID); err == nil { + t.Fatal("another user revoked a binding they do not own") + } + if _, err := db.ResolveRobotRBACAccess("wecom", "t:corp|u:one"); err != nil { + t.Fatalf("unauthorized revocation changed binding: %v", err) + } + if err := db.DeleteRobotUserBindingForUser(bindings[0].ID, u1.ID); err != nil { + t.Fatal(err) + } + if _, err := db.ResolveRobotRBACAccess("wecom", "t:corp|u:one"); err == nil { + t.Fatal("revoked binding still resolves") + } +}