diff --git a/internal/handler/agent.go b/internal/handler/agent.go index 9c27fb8c..ac2ea714 100644 --- a/internal/handler/agent.go +++ b/internal/handler/agent.go @@ -24,6 +24,7 @@ import ( "cyberstrike-ai/internal/multiagent" "cyberstrike-ai/internal/openai" "cyberstrike-ai/internal/reasoning" + "cyberstrike-ai/internal/security" "github.com/gin-gonic/gin" "github.com/robfig/cron/v3" @@ -844,6 +845,36 @@ func (h *AgentHandler) publishProgressToTaskEventBus(conversationID, eventType, h.taskEventBus.Publish(conversationID, sseLine) } +// enrichProgressEventData 为 SSE / taskEventBus 事件补齐 conversationId、messageId，便于前端懒加载过程详情。 +func enrichProgressEventData(data interface{}, conversationID, assistantMessageID string) interface{} { + if strings.TrimSpace(conversationID) == "" && strings.TrimSpace(assistantMessageID) == "" { + return data + } + var m map[string]interface{} + switch v := data.(type) { + case map[string]interface{}: + m = make(map[string]interface{}, len(v)+2) + for k, val := range v { + m[k] = val + } + case nil: + m = make(map[string]interface{}, 2) + default: + m = map[string]interface{}{"payload": data} + } + if id := strings.TrimSpace(assistantMessageID); id != "" { + if existing, ok := m["messageId"]; !ok || strings.TrimSpace(fmt.Sprint(existing)) == "" { + m["messageId"] = id + } + } + if id := strings.TrimSpace(conversationID); id != "" { + if existing, ok := m["conversationId"]; !ok || strings.TrimSpace(fmt.Sprint(existing)) == "" { + m["conversationId"] = id + } + } + return m +} + // createProgressCallback 创建进度回调函数，用于保存processDetails // sendEventFunc: 可选的流式事件发送函数，如果为nil则不发送流式事件 func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun context.CancelCauseFunc, conversationID, assistantMessageID string, sendEventFunc func(eventType, message string, data interface{})) agent.ProgressCallback { @@ -997,10 +1028,11 @@ func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun // 流式：写 HTTP SSE；非流式（机器人等）：镜像到 taskEventBus 供 Web 订阅。 // 工具事件需先落库拿 processDetailId，再向前端发送摘要，避免大 payload 默认进入浏览器。 if !deferToolProgressSend { + clientData := enrichProgressEventData(data, conversationID, assistantMessageID) if sendEventFunc != nil { - sendEventFunc(eventType, message, data) + sendEventFunc(eventType, message, clientData) } else { - h.publishProgressToTaskEventBus(conversationID, eventType, message, data) + h.publishProgressToTaskEventBus(conversationID, eventType, message, clientData) } } @@ -1374,7 +1406,7 @@ func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun h.logger.Warn("保存过程详情失败", zap.Error(err), zap.String("eventType", eventType)) } if deferToolProgressSend { - clientData := summarizeProcessDetailData(eventType, data) + clientData := enrichProgressEventData(summarizeProcessDetailData(eventType, data), conversationID, assistantMessageID) if m, ok := clientData.(map[string]interface{}); ok { m["processDetailId"] = processDetailID } @@ -1385,7 +1417,7 @@ func (h *AgentHandler) createProgressCallback(runCtx context.Context, cancelRun } } } else if deferToolProgressSend { - clientData := summarizeProcessDetailData(eventType, data) + clientData := enrichProgressEventData(summarizeProcessDetailData(eventType, data), conversationID, assistantMessageID) if sendEventFunc != nil { sendEventFunc(eventType, message, clientData) } else { @@ -1682,6 +1714,12 @@ func (h *AgentHandler) CreateBatchQueue(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "没有有效的任务"}) return } + if session, ok := security.CurrentSession(c); ok && h.db != nil && session.Scope != database.RBACScopeAll && strings.TrimSpace(req.ProjectID) != "" { + if !h.db.UserCanAccessResource(session.UserID, session.Scope, "project", strings.TrimSpace(req.ProjectID)) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权在该项目下创建批量任务"}) + return + } + } agentMode := config.NormalizeAgentMode(req.AgentMode) scheduleMode := normalizeBatchQueueScheduleMode(req.ScheduleMode) @@ -1706,6 +1744,10 @@ func (h *AgentHandler) CreateBatchQueue(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": createErr.Error()}) return } + if session, ok := security.CurrentSession(c); ok && h.db != nil { + _ = h.db.SetResourceOwner("batch_task", queue.ID, session.UserID) + _ = h.db.AssignResourceToUser(session.UserID, "batch_task", queue.ID) + } started := false if req.ExecuteNow { ok, err := h.startBatchQueueExecution(queue.ID, false) @@ -1793,7 +1835,8 @@ func (h *AgentHandler) ListBatchQueues(c *gin.Context) { } // 获取队列列表和总数 - queues, total, err := h.batchTaskManager.ListQueues(limit, offset, status, keyword) + session, _ := security.CurrentSession(c) + queues, total, err := h.batchTaskManager.ListQueuesForAccess(limit, offset, status, keyword, session.UserID, session.Scope) if err != nil { h.logger.Error("获取批量任务队列列表失败", zap.Error(err)) c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) @@ -2261,6 +2304,7 @@ func (h *AgentHandler) loadHistoryFromAgentTrace(conversationID string) ([]agent } messageCount := len(messagesArray) + modelFacingTrace := agent.IsModelFacingTraceJSON(traceInputJSON) h.logger.Info("使用保存的代理轨迹恢复历史上下文", zap.String("conversationId", conversationID), @@ -2275,6 +2319,7 @@ func (h *AgentHandler) loadHistoryFromAgentTrace(conversationID string) ([]agent agentMessages := make([]agent.ChatMessage, 0, len(messagesArray)) for _, msgMap := range messagesArray { msg := agent.ChatMessage{} + msg.ModelFacingTrace = modelFacingTrace // 解析role if role, ok := msgMap["role"].(string); ok { diff --git a/internal/handler/agent_progress_callback_test.go b/internal/handler/agent_progress_callback_test.go index 6eb13e31..5447c9a7 100644 --- a/internal/handler/agent_progress_callback_test.go +++ b/internal/handler/agent_progress_callback_test.go @@ -97,3 +97,26 @@ func TestCreateProgressCallback_FlushesReasoningOnDone(t *testing.T) { t.Fatalf("expected reasoning_chain persisted on done, got %+v", details) } } + +func TestEnrichProgressEventData(t *testing.T) { + t.Run("fills ids", func(t *testing.T) { + out := enrichProgressEventData(map[string]interface{}{"source": "eino"}, "conv-1", "msg-1") + m, ok := out.(map[string]interface{}) + if !ok { + t.Fatalf("expected map, got %T", out) + } + if m["conversationId"] != "conv-1" || m["messageId"] != "msg-1" { + t.Fatalf("unexpected enrichment: %+v", m) + } + }) + t.Run("preserves existing ids", func(t *testing.T) { + out := enrichProgressEventData(map[string]interface{}{ + "conversationId": "keep-conv", + "messageId": "keep-msg", + }, "conv-1", "msg-1") + m := out.(map[string]interface{}) + if m["conversationId"] != "keep-conv" || m["messageId"] != "keep-msg" { + t.Fatalf("should not overwrite existing ids: %+v", m) + } + }) +} diff --git a/internal/handler/auth.go b/internal/handler/auth.go index a0e940d2..c36e72e2 100644 --- a/internal/handler/auth.go +++ b/internal/handler/auth.go @@ -38,6 +38,7 @@ func NewAuthHandler(manager *security.AuthManager, cfg *config.Config, configPat } type loginRequest struct { + Username string `json:"username"` Password string `json:"password" binding:"required"` } @@ -54,7 +55,7 @@ func (h *AuthHandler) Login(c *gin.Context) { return } - token, expiresAt, err := h.manager.Authenticate(req.Password) + token, expiresAt, err := h.manager.Authenticate(req.Username, req.Password) if err != nil { if h.audit != nil { h.audit.Record(c, audit.Entry{ @@ -68,6 +69,7 @@ func (h *AuthHandler) Login(c *gin.Context) { c.JSON(http.StatusUnauthorized, gin.H{"error": "密码错误"}) return } + session, _ := h.manager.ValidateToken(token) if h.audit != nil { h.audit.Record(c, audit.Entry{ @@ -86,6 +88,14 @@ func (h *AuthHandler) Login(c *gin.Context) { "token": token, "expires_at": expiresAt.UTC().Format(time.RFC3339), "session_duration_hr": h.manager.SessionDurationHours(), + "user": gin.H{ + "id": session.UserID, + "username": session.Username, + "display_name": session.DisplayName, + }, + "roles": session.Roles, + "permissions": permissionKeys(session.Permissions), + "scope": session.Scope, }) } @@ -139,7 +149,11 @@ func (h *AuthHandler) ChangePassword(c *gin.Context) { return } - if !h.manager.CheckPassword(oldPassword) { + session, _ := security.CurrentSession(c) + if session.Username == "" { + session.Username = "admin" + } + if !h.manager.CheckUserPassword(session.Username, oldPassword) { if h.audit != nil { h.audit.Record(c, audit.Entry{ Level: "warn", @@ -153,27 +167,35 @@ func (h *AuthHandler) ChangePassword(c *gin.Context) { return } - if err := config.PersistAuthPassword(h.configPath, newPassword); err != nil { - if h.logger != nil { - h.logger.Error("保存新密码失败", zap.Error(err)) + if session.UserID == "" || session.UserID == "admin" { + if err := config.PersistAuthPassword(h.configPath, newPassword); err != nil { + if h.logger != nil { + h.logger.Error("保存新密码失败", zap.Error(err)) + } + c.JSON(http.StatusInternalServerError, gin.H{"error": "保存新密码失败，请重试"}) + return } - c.JSON(http.StatusInternalServerError, gin.H{"error": "保存新密码失败，请重试"}) + + if err := h.manager.UpdateConfig(newPassword, h.config.Auth.SessionDurationHours); err != nil { + if h.logger != nil { + h.logger.Error("更新认证配置失败", zap.Error(err)) + } + c.JSON(http.StatusInternalServerError, gin.H{"error": "更新认证配置失败"}) + return + } + + h.config.Auth.Password = newPassword + h.config.Auth.GeneratedPassword = "" + h.config.Auth.GeneratedPasswordPersisted = false + h.config.Auth.GeneratedPasswordPersistErr = "" + } else if err := h.manager.UpdateUserPassword(session.UserID, newPassword); err != nil { + if h.logger != nil { + h.logger.Error("更新用户密码失败", zap.Error(err)) + } + c.JSON(http.StatusInternalServerError, gin.H{"error": "更新用户密码失败"}) return } - if err := h.manager.UpdateConfig(newPassword, h.config.Auth.SessionDurationHours); err != nil { - if h.logger != nil { - h.logger.Error("更新认证配置失败", zap.Error(err)) - } - c.JSON(http.StatusInternalServerError, gin.H{"error": "更新认证配置失败"}) - return - } - - h.config.Auth.Password = newPassword - h.config.Auth.GeneratedPassword = "" - h.config.Auth.GeneratedPasswordPersisted = false - h.config.Auth.GeneratedPasswordPersistErr = "" - if h.logger != nil { h.logger.Info("登录密码已更新，所有会话已失效") } @@ -207,5 +229,23 @@ func (h *AuthHandler) Validate(c *gin.Context) { c.JSON(http.StatusOK, gin.H{ "token": session.Token, "expires_at": session.ExpiresAt.UTC().Format(time.RFC3339), + "user": gin.H{ + "id": session.UserID, + "username": session.Username, + "display_name": session.DisplayName, + }, + "roles": session.Roles, + "permissions": permissionKeys(session.Permissions), + "scope": session.Scope, }) } + +func permissionKeys(perms map[string]bool) []string { + keys := make([]string, 0, len(perms)) + for key, ok := range perms { + if ok { + keys = append(keys, key) + } + } + return keys +} diff --git a/internal/handler/batch_task_manager.go b/internal/handler/batch_task_manager.go index b99822d0..9db9f266 100644 --- a/internal/handler/batch_task_manager.go +++ b/internal/handler/batch_task_manager.go @@ -98,8 +98,8 @@ type BatchTaskManager struct { logger *zap.Logger queues map[string]*BatchTaskQueue taskCancels map[string]map[string]context.CancelFunc // queueID -> taskID -> 取消函数 - singleRunTasks map[string]string // queueID -> taskID，单条执行完成后暂停队列 - queueExecutors map[string]struct{} // executeBatchQueue 协程活跃标记（与队列 status 解耦） + singleRunTasks map[string]string // queueID -> taskID，单条执行完成后暂停队列 + queueExecutors map[string]struct{} // executeBatchQueue 协程活跃标记（与队列 status 解耦） mu sync.RWMutex } @@ -426,20 +426,24 @@ func (m *BatchTaskManager) GetAllQueues() []*BatchTaskQueue { // ListQueues 列出队列（支持筛选和分页） func (m *BatchTaskManager) ListQueues(limit, offset int, status, keyword string) ([]*BatchTaskQueue, int, error) { + return m.ListQueuesForAccess(limit, offset, status, keyword, "", "") +} + +func (m *BatchTaskManager) ListQueuesForAccess(limit, offset int, status, keyword, userID, scope string) ([]*BatchTaskQueue, int, error) { var queues []*BatchTaskQueue var total int // 如果数据库可用，从数据库查询 if m.db != nil { // 获取总数 - count, err := m.db.CountBatchQueues(status, keyword) + count, err := m.db.CountBatchQueuesForAccess(status, keyword, userID, scope) if err != nil { return nil, 0, fmt.Errorf("统计队列总数失败: %w", err) } total = count // 获取队列列表（只获取ID） - queueRows, err := m.db.ListBatchQueues(limit, offset, status, keyword) + queueRows, err := m.db.ListBatchQueuesForAccess(limit, offset, status, keyword, userID, scope) if err != nil { return nil, 0, fmt.Errorf("查询队列列表失败: %w", err) } diff --git a/internal/handler/c2.go b/internal/handler/c2.go index feebe888..fe977e30 100644 --- a/internal/handler/c2.go +++ b/internal/handler/c2.go @@ -17,6 +17,7 @@ import ( "cyberstrike-ai/internal/audit" "cyberstrike-ai/internal/c2" "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" "github.com/gin-gonic/gin" "github.com/google/uuid" @@ -59,7 +60,7 @@ func (h *C2Handler) SetManager(m *c2.Manager) { // ListListeners 获取监听器列表 func (h *C2Handler) ListListeners(c *gin.Context) { - listeners, err := h.mgr().DB().ListC2Listeners() + listeners, err := h.mgr().DB().ListC2ListenersForAccess(c2AccessFromContext(c)) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return @@ -109,6 +110,11 @@ func (h *C2Handler) CreateListener(c *gin.Context) { c.JSON(code, gin.H{"error": err.Error()}) return } + if session, ok := security.CurrentSession(c); ok { + listener.OwnerUserID = session.UserID + _ = h.mgr().DB().SetResourceOwner("c2_listener", listener.ID, session.UserID) + _ = h.mgr().DB().AssignResourceToUser(session.UserID, "c2_listener", listener.ID) + } implantToken := listener.ImplantToken listener.EncryptionKey = "" listener.ImplantToken = "" @@ -282,7 +288,7 @@ func (h *C2Handler) ListSessions(c *gin.Context) { filter.Suspicious = true } - sessions, err := h.mgr().DB().ListC2Sessions(filter) + sessions, err := h.mgr().DB().ListC2SessionsForAccess(filter, c2AccessFromContext(c)) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return @@ -304,10 +310,10 @@ func (h *C2Handler) GetSession(c *gin.Context) { } // 获取最近任务 - tasks, _ := h.mgr().DB().ListC2Tasks(database.ListC2TasksFilter{ + tasks, _ := h.mgr().DB().ListC2TasksForAccess(database.ListC2TasksFilter{ SessionID: id, Limit: 20, - }) + }, c2AccessFromContext(c)) c.JSON(http.StatusOK, gin.H{ "session": session, @@ -341,7 +347,7 @@ func (h *C2Handler) DeleteSessions(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "ids is required"}) return } - n, err := h.mgr().DB().DeleteC2SessionsByIDs(req.IDs) + n, err := h.mgr().DB().DeleteC2SessionsByIDsForAccess(req.IDs, c2AccessFromContext(c)) if err != nil { if errors.Is(err, database.ErrNoValidC2SessionIDs) { c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) @@ -433,24 +439,25 @@ func (h *C2Handler) ListTasks(c *gin.Context) { } } - tasks, err := h.mgr().DB().ListC2Tasks(filter) + access := c2AccessFromContext(c) + tasks, err := h.mgr().DB().ListC2TasksForAccess(filter, access) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return } // 仪表盘「待审任务」为全局 queued/pending 数量，与列表 session 过滤无关 - pendingN, _ := h.mgr().DB().CountC2TasksQueuedOrPending("") + pendingN, _ := h.mgr().DB().CountC2TasksQueuedOrPendingForAccess("", access) if !paginated { c.JSON(http.StatusOK, gin.H{ - "tasks": tasks, - "pending_queued_count": pendingN, + "tasks": tasks, + "pending_queued_count": pendingN, }) return } - total, err := h.mgr().DB().CountC2Tasks(filter) + total, err := h.mgr().DB().CountC2TasksForAccess(filter, access) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return @@ -477,7 +484,7 @@ func (h *C2Handler) DeleteTasks(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "ids is required"}) return } - n, err := h.mgr().DB().DeleteC2TasksByIDs(req.IDs) + n, err := h.mgr().DB().DeleteC2TasksByIDsForAccess(req.IDs, c2AccessFromContext(c)) if err != nil { if errors.Is(err, database.ErrNoValidC2TaskIDs) { c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) @@ -522,6 +529,13 @@ func (h *C2Handler) CreateTask(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) return } + if strings.TrimSpace(req.SessionID) == "" { + req.SessionID = strings.TrimSpace(c.Param("id")) + } + if !h.c2ResourceAllowed(c, "c2_session", req.SessionID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } input := c2.EnqueueTaskInput{ SessionID: req.SessionID, @@ -621,6 +635,10 @@ func (h *C2Handler) PayloadOneliner(c *gin.Context) { c.JSON(http.StatusNotFound, gin.H{"error": "listener not found"}) return } + if !h.c2ResourceAllowed(c, "c2_listener", req.ListenerID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } host := c2.ResolveBeaconDialHost(listener, strings.TrimSpace(req.Host), h.logger, listener.ID) @@ -684,6 +702,10 @@ func (h *C2Handler) PayloadBuild(c *gin.Context) { c.JSON(http.StatusNotFound, gin.H{"error": "listener not found"}) return } + if !h.c2ResourceAllowed(c, "c2_listener", req.ListenerID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } builder := c2.NewPayloadBuilder(h.mgr(), h.logger, "", "") input := c2.PayloadBuilderInput{ @@ -779,7 +801,8 @@ func (h *C2Handler) ListEvents(c *gin.Context) { } } - events, err := h.mgr().DB().ListC2Events(filter) + access := c2AccessFromContext(c) + events, err := h.mgr().DB().ListC2EventsForAccess(filter, access) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return @@ -788,7 +811,7 @@ func (h *C2Handler) ListEvents(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"events": events}) return } - total, err := h.mgr().DB().CountC2Events(filter) + total, err := h.mgr().DB().CountC2EventsForAccess(filter, access) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return @@ -814,7 +837,7 @@ func (h *C2Handler) DeleteEvents(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "ids is required"}) return } - n, err := h.mgr().DB().DeleteC2EventsByIDs(req.IDs) + n, err := h.mgr().DB().DeleteC2EventsByIDsForAccess(req.IDs, c2AccessFromContext(c)) if err != nil { if errors.Is(err, database.ErrNoValidC2EventIDs) { c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) @@ -851,6 +874,9 @@ func (h *C2Handler) EventStream(c *gin.Context) { if !ok { return false } + if !h.c2EventAllowed(c, e) { + return true + } data, _ := json.Marshal(e) fmt.Fprintf(w, "data: %s\n\n", data) return true @@ -964,6 +990,10 @@ func (h *C2Handler) UploadFileForImplant(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "session_id and remote_path required"}) return } + if !h.c2ResourceAllowed(c, "c2_session", sessionID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } file, header, err := c.Request.FormFile("file") if err != nil { @@ -1018,6 +1048,10 @@ func (h *C2Handler) ListFiles(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "session_id required"}) return } + if !h.c2ResourceAllowed(c, "c2_session", sessionID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } files, err := h.mgr().DB().ListC2FilesBySession(sessionID) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) @@ -1038,6 +1072,10 @@ func (h *C2Handler) DownloadResultFile(c *gin.Context) { c.JSON(http.StatusNotFound, gin.H{"error": "task not found"}) return } + if !h.c2ResourceAllowed(c, "c2_task", taskID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } if task.ResultBlobPath == "" { c.JSON(http.StatusNotFound, gin.H{"error": "no result file for this task"}) return @@ -1053,6 +1091,42 @@ func osCreate(path string) (*os.File, error) { return os.Create(path) } +func c2AccessFromContext(c *gin.Context) database.RBACListAccess { + session, ok := security.CurrentSession(c) + if !ok { + return database.RBACListAccess{} + } + return database.RBACListAccess{UserID: session.UserID, Scope: session.Scope} +} + +func (h *C2Handler) c2ResourceAllowed(c *gin.Context, resourceType, resourceID string) bool { + session, ok := security.CurrentSession(c) + if !ok { + return false + } + return h.mgr().DB().UserCanAccessResource(session.UserID, session.Scope, resourceType, resourceID) +} + +func (h *C2Handler) c2EventAllowed(c *gin.Context, e *c2.Event) bool { + if e == nil { + return false + } + session, ok := security.CurrentSession(c) + if !ok { + return false + } + if session.Scope == database.RBACScopeAll { + return true + } + if strings.TrimSpace(e.SessionID) != "" { + return h.mgr().DB().UserCanAccessResource(session.UserID, session.Scope, "c2_session", e.SessionID) + } + if strings.TrimSpace(e.TaskID) != "" { + return h.mgr().DB().UserCanAccessResource(session.UserID, session.Scope, "c2_task", e.TaskID) + } + return false +} + // ============================================================================ // 辅助函数（firstNonEmpty 已在 vulnerability.go 中定义） // ============================================================================ diff --git a/internal/handler/conversation.go b/internal/handler/conversation.go index f98db59e..3505ae2a 100644 --- a/internal/handler/conversation.go +++ b/internal/handler/conversation.go @@ -8,6 +8,7 @@ import ( "cyberstrike-ai/internal/audit" "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" "github.com/gin-gonic/gin" "go.uber.org/zap" ) @@ -69,12 +70,23 @@ func (h *ConversationHandler) CreateConversation(c *gin.Context) { meta := audit.ConversationCreateMetaFromGin(c, "api") meta.ProjectID = strings.TrimSpace(req.ProjectID) + if !h.conversationProjectAllowed(c, meta.ProjectID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问目标项目"}) + return + } conv, err := h.db.CreateConversation(title, meta) if err != nil { h.logger.Error("创建对话失败", zap.Error(err)) c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return } + if session, ok := security.CurrentSession(c); ok { + _ = h.db.SetResourceOwner("conversation", conv.ID, session.UserID) + _ = h.db.AssignResourceToUser(session.UserID, "conversation", conv.ID) + if conv.ProjectID != "" { + _ = h.db.AssignResourceToUser(session.UserID, "project", conv.ProjectID) + } + } c.JSON(http.StatusOK, conv) } @@ -91,11 +103,28 @@ func (h *ConversationHandler) SetConversationProject(c *gin.Context) { c.JSON(http.StatusNotFound, gin.H{"error": "对话不存在"}) return } - if err := h.db.SetConversationProjectID(id, req.ProjectID); err != nil { + projectID := strings.TrimSpace(req.ProjectID) + if !h.conversationProjectAllowed(c, projectID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问目标项目"}) + return + } + if err := h.db.SetConversationProjectID(id, projectID); err != nil { c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) return } - c.JSON(http.StatusOK, gin.H{"success": true, "projectId": strings.TrimSpace(req.ProjectID)}) + c.JSON(http.StatusOK, gin.H{"success": true, "projectId": projectID}) +} + +func (h *ConversationHandler) conversationProjectAllowed(c *gin.Context, projectID string) bool { + projectID = strings.TrimSpace(projectID) + if projectID == "" { + return true + } + session, ok := security.CurrentSession(c) + if !ok { + return false + } + return h.db.UserCanAccessResource(session.UserID, session.Scope, "project", projectID) } // ListConversations 列出对话 @@ -118,19 +147,20 @@ func (h *ConversationHandler) ListConversations(c *gin.Context) { excludeGrouped := strings.TrimSpace(search) == "" && projectID == "" && (c.Query("exclude_grouped") == "true" || c.Query("exclude_grouped") == "1") sortBy := strings.TrimSpace(c.Query("sort_by")) + session, _ := security.CurrentSession(c) var conversations []*database.Conversation var total int var err error if excludeGrouped { - conversations, err = h.db.ListUngroupedConversations(limit, offset, sortBy, projectID) + conversations, err = h.db.ListUngroupedConversationsForAccess(limit, offset, sortBy, projectID, session.UserID, session.Scope) if err == nil { - total, err = h.db.CountUngroupedConversations(projectID) + total, err = h.db.CountUngroupedConversationsForAccess(projectID, session.UserID, session.Scope) } } else { - conversations, err = h.db.ListConversations(limit, offset, search, sortBy, projectID) + conversations, err = h.db.ListConversationsForAccess(limit, offset, search, sortBy, projectID, session.UserID, session.Scope) if err == nil { - total, err = h.db.CountConversations(search, projectID) + total, err = h.db.CountConversationsForAccess(search, projectID, session.UserID, session.Scope) } } if err != nil { @@ -266,8 +296,20 @@ func (h *ConversationHandler) GetMessageProcessDetails(c *gin.Context) { } details = database.DedupeConsecutiveProcessDetails(details) out := processDetailsToJSON(h.logger, details, false) + // A page may end between tool_call and tool_result. Return the full-history + // execution summary so the UI can render terminal status without pretending + // that an unloaded result is still running. + summary, summaryErr := h.db.GetProcessDetailsSummary(messageID) + if summaryErr != nil { + h.logger.Warn("获取分页工具执行状态失败", zap.Error(summaryErr), zap.String("messageID", messageID)) + } + var toolExecutions []database.ProcessDetailsToolExecution + if summary != nil { + toolExecutions = summary.ToolExecutions + } c.JSON(http.StatusOK, gin.H{ "processDetails": out, + "toolExecutions": toolExecutions, "total": total, "offset": offset, "limit": limit, diff --git a/internal/handler/conversation_process_details_test.go b/internal/handler/conversation_process_details_test.go new file mode 100644 index 00000000..8c3ba700 --- /dev/null +++ b/internal/handler/conversation_process_details_test.go @@ -0,0 +1,75 @@ +package handler + +import ( + "encoding/json" + "fmt" + "net/http/httptest" + "path/filepath" + "testing" + + "cyberstrike-ai/internal/database" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +func TestProcessDetailsPageIncludesTerminalToolStatusAcrossPageBoundary(t *testing.T) { + gin.SetMode(gin.TestMode) + db, err := database.NewDB(filepath.Join(t.TempDir(), "process-details-page.db"), zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + conversation, err := db.CreateConversation("page boundary", database.ConversationCreateMeta{}) + if err != nil { + t.Fatalf("CreateConversation: %v", err) + } + message, err := db.AddMessage(conversation.ID, "assistant", "done", nil) + if err != nil { + t.Fatalf("AddMessage: %v", err) + } + for i := 1; i <= 4; i++ { + id := fmt.Sprintf("call-%d", i) + if err := db.AddProcessDetail(message.ID, conversation.ID, "tool_call", "call", map[string]interface{}{ + "toolName": "http-framework-test", "toolCallId": id, "index": i, "total": 4, + }); err != nil { + t.Fatalf("AddProcessDetail(tool_call): %v", err) + } + } + for i := 1; i <= 4; i++ { + id := fmt.Sprintf("call-%d", i) + if err := db.AddProcessDetail(message.ID, conversation.ID, "tool_result", "result", map[string]interface{}{ + "toolName": "http-framework-test", "toolCallId": id, "success": true, + }); err != nil { + t.Fatalf("AddProcessDetail(tool_result): %v", err) + } + } + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/api/messages/"+message.ID+"/process-details?limit=6&offset=0", nil) + c.Params = gin.Params{{Key: "id", Value: message.ID}} + NewConversationHandler(db, zap.NewNop()).GetMessageProcessDetails(c) + if w.Code != 200 { + t.Fatalf("status = %d: %s", w.Code, w.Body.String()) + } + var response struct { + HasMore bool `json:"hasMore"` + ProcessDetails []map[string]interface{} `json:"processDetails"` + ToolExecutions []database.ProcessDetailsToolExecution `json:"toolExecutions"` + } + if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil { + t.Fatalf("decode response: %v", err) + } + if !response.HasMore || len(response.ProcessDetails) != 6 { + t.Fatalf("page hasMore=%v details=%d, want true/6", response.HasMore, len(response.ProcessDetails)) + } + if len(response.ToolExecutions) != 4 { + t.Fatalf("tool executions = %d, want 4", len(response.ToolExecutions)) + } + for i, execution := range response.ToolExecutions { + if execution.Status != "completed" { + t.Fatalf("execution %d status = %q, want completed", i, execution.Status) + } + } +} diff --git a/internal/handler/conversation_rbac_test.go b/internal/handler/conversation_rbac_test.go new file mode 100644 index 00000000..ea560e21 --- /dev/null +++ b/internal/handler/conversation_rbac_test.go @@ -0,0 +1,118 @@ +package handler + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "path/filepath" + "testing" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +func TestCreateConversationRequiresProjectAccess(t *testing.T) { + gin.SetMode(gin.TestMode) + db, user := setupConversationRBACTest(t) + project, err := db.CreateProject(&database.Project{Name: "hidden"}) + if err != nil { + t.Fatalf("CreateProject: %v", err) + } + handler := NewConversationHandler(db, zap.NewNop()) + + w := performConversationRequest(user, http.MethodPost, "/api/conversations", map[string]string{ + "title": "blocked", + "projectId": project.ID, + }, handler.CreateConversation) + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String()) + } + + if err := db.AssignResourceToUser(user.ID, "project", project.ID); err != nil { + t.Fatalf("AssignResourceToUser: %v", err) + } + w = performConversationRequest(user, http.MethodPost, "/api/conversations", map[string]string{ + "title": "allowed", + "projectId": project.ID, + }, handler.CreateConversation) + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusOK, w.Body.String()) + } +} + +func TestSetConversationProjectRequiresProjectAccess(t *testing.T) { + gin.SetMode(gin.TestMode) + db, user := setupConversationRBACTest(t) + project, err := db.CreateProject(&database.Project{Name: "hidden"}) + if err != nil { + t.Fatalf("CreateProject: %v", err) + } + conv, err := db.CreateConversation("owned", database.ConversationCreateMeta{}) + if err != nil { + t.Fatalf("CreateConversation: %v", err) + } + if err := db.SetResourceOwner("conversation", conv.ID, user.ID); err != nil { + t.Fatalf("SetResourceOwner: %v", err) + } + if err := db.AssignResourceToUser(user.ID, "conversation", conv.ID); err != nil { + t.Fatalf("AssignResourceToUser conversation: %v", err) + } + handler := NewConversationHandler(db, zap.NewNop()) + + w := performConversationRequest(user, http.MethodPut, "/api/conversations/"+conv.ID+"/project", map[string]string{ + "projectId": project.ID, + }, func(c *gin.Context) { + c.Params = gin.Params{{Key: "id", Value: conv.ID}} + handler.SetConversationProject(c) + }) + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String()) + } + + if err := db.AssignResourceToUser(user.ID, "project", project.ID); err != nil { + t.Fatalf("AssignResourceToUser project: %v", err) + } + w = performConversationRequest(user, http.MethodPut, "/api/conversations/"+conv.ID+"/project", map[string]string{ + "projectId": project.ID, + }, func(c *gin.Context) { + c.Params = gin.Params{{Key: "id", Value: conv.ID}} + handler.SetConversationProject(c) + }) + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusOK, w.Body.String()) + } +} + +func setupConversationRBACTest(t *testing.T) (*database.DB, *database.RBACUser) { + t.Helper() + db, err := database.NewDB(filepath.Join(t.TempDir(), "conversation-rbac.db"), zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + user, err := db.CreateRBACUser("operator1", "Operator One", "hash", true, nil) + if err != nil { + t.Fatalf("CreateRBACUser: %v", err) + } + return db, user +} + +func performConversationRequest(user *database.RBACUser, method, path string, body map[string]string, handler gin.HandlerFunc) *httptest.ResponseRecorder { + payload, _ := json.Marshal(body) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(method, path, bytes.NewReader(payload)) + c.Request.Header.Set("Content-Type", "application/json") + c.Set(security.ContextSessionKey, security.Session{ + UserID: user.ID, + Username: user.Username, + Permissions: map[string]bool{"chat:write": true}, + Scope: database.RBACScopeAssigned, + }) + handler(c) + return w +} diff --git a/internal/handler/group.go b/internal/handler/group.go index 495e7695..db44e228 100644 --- a/internal/handler/group.go +++ b/internal/handler/group.go @@ -5,6 +5,7 @@ import ( "time" "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" "github.com/gin-gonic/gin" "go.uber.org/zap" @@ -152,6 +153,10 @@ func (h *GroupHandler) AddConversationToGroup(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) return } + if !h.groupConversationAllowed(c, req.ConversationID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } if err := h.db.AddConversationToGroup(req.ConversationID, req.GroupID); err != nil { h.logger.Error("添加对话到分组失败", zap.Error(err)) @@ -166,6 +171,10 @@ func (h *GroupHandler) AddConversationToGroup(c *gin.Context) { func (h *GroupHandler) RemoveConversationFromGroup(c *gin.Context) { conversationID := c.Param("conversationId") groupID := c.Param("id") + if !h.groupConversationAllowed(c, conversationID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } if err := h.db.RemoveConversationFromGroup(conversationID, groupID); err != nil { h.logger.Error("从分组中移除对话失败", zap.Error(err)) @@ -210,6 +219,9 @@ func (h *GroupHandler) GetGroupConversations(c *gin.Context) { // 获取每个对话在分组中的置顶状态 groupConvs := make([]GroupConversation, 0, len(conversations)) for _, conv := range conversations { + if conv == nil || !h.groupConversationAllowed(c, conv.ID) { + continue + } // 查询分组内置顶状态 var groupPinned int err := h.db.QueryRow( @@ -242,8 +254,14 @@ func (h *GroupHandler) GetAllMappings(c *gin.Context) { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return } + filtered := mappings[:0] + for _, mapping := range mappings { + if h.groupConversationAllowed(c, mapping.ConversationID) { + filtered = append(filtered, mapping) + } + } - c.JSON(http.StatusOK, mappings) + c.JSON(http.StatusOK, filtered) } // UpdateConversationPinnedRequest 更新对话置顶状态请求 @@ -254,6 +272,10 @@ type UpdateConversationPinnedRequest struct { // UpdateConversationPinned 更新对话置顶状态 func (h *GroupHandler) UpdateConversationPinned(c *gin.Context) { conversationID := c.Param("id") + if !h.groupConversationAllowed(c, conversationID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } var req UpdateConversationPinnedRequest if err := c.ShouldBindJSON(&req); err != nil { @@ -303,6 +325,10 @@ type UpdateConversationPinnedInGroupRequest struct { func (h *GroupHandler) UpdateConversationPinnedInGroup(c *gin.Context) { groupID := c.Param("id") conversationID := c.Param("conversationId") + if !h.groupConversationAllowed(c, conversationID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } var req UpdateConversationPinnedInGroupRequest if err := c.ShouldBindJSON(&req); err != nil { @@ -318,3 +344,11 @@ func (h *GroupHandler) UpdateConversationPinnedInGroup(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"message": "更新成功"}) } + +func (h *GroupHandler) groupConversationAllowed(c *gin.Context, conversationID string) bool { + session, ok := security.CurrentSession(c) + if !ok { + return false + } + return h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", conversationID) +} diff --git a/internal/handler/hitl.go b/internal/handler/hitl.go index 1a0f7146..0fb5f356 100644 --- a/internal/handler/hitl.go +++ b/internal/handler/hitl.go @@ -611,6 +611,7 @@ func (h *AgentHandler) ListHITLPending(c *gin.Context) { offset := (page - 1) * pageSize q, args := h.buildHitlListQuery(false) q, args = h.appendHitlListFilters(q, args, c) + q, args = appendConversationAccessSQL(q, args, "conversation_id", notificationAccessFromContext(c)) total, err := h.countHitlQuery(q, args) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) @@ -649,6 +650,10 @@ func (h *AgentHandler) DecideHITLInterrupt(c *gin.Context) { c.JSON(500, gin.H{"error": "hitl manager unavailable"}) return } + if !h.hitlInterruptAllowed(c, req.InterruptID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } if err := h.hitlManager.ResolveInterrupt(req.InterruptID, req.Decision, req.Comment, req.EditedArguments); err != nil { c.JSON(http.StatusConflict, gin.H{"error": err.Error()}) return @@ -673,6 +678,10 @@ func (h *AgentHandler) DismissHITLInterrupt(c *gin.Context) { c.JSON(500, gin.H{"error": "hitl manager unavailable"}) return } + if !h.hitlInterruptAllowed(c, req.InterruptID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } res, err := h.db.Exec(`UPDATE hitl_interrupts SET status='cancelled', decision='reject', decision_comment='dismissed by user', decided_at=CURRENT_TIMESTAMP, decided_by='human' WHERE id=? AND status='pending'`, req.InterruptID) @@ -728,7 +737,6 @@ func (h *AgentHandler) interceptHITLForEinoTool(runCtx context.Context, cancelRu return arguments, nil } - type hitlConfigReq struct { ConversationID string `json:"conversationId" binding:"required"` HITLRequest @@ -740,6 +748,10 @@ func (h *AgentHandler) GetHITLConversationConfig(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "conversationId is required"}) return } + if !h.hitlConversationAllowed(c, conversationID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } cfg, err := h.loadHITLConversationConfig(conversationID) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) @@ -770,6 +782,10 @@ func (h *AgentHandler) UpsertHITLConversationConfig(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) return } + if !h.hitlConversationAllowed(c, req.ConversationID) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } req.Mode = normalizeHitlMode(req.Mode) req.Reviewer = normalizeHitlReviewer(req.Reviewer) if strings.TrimSpace(req.Reviewer) == "" { @@ -868,8 +884,8 @@ func (h *AgentHandler) SetHITLGlobalToolWhitelist(c *gin.Context) { h.audit.RecordOK(c, "hitl", "tool_whitelist_update", "HITL 全局白名单更新", "hitl_config", "tool_whitelist", nil) } c.JSON(http.StatusOK, gin.H{ - "ok": true, - "toolWhitelist": h.hitlConfigGlobalToolWhitelist(), + "ok": true, + "toolWhitelist": h.hitlConfigGlobalToolWhitelist(), "hitlGlobalToolWhitelist": h.hitlConfigGlobalToolWhitelist(), "hitlGlobalWhitelistMerged": false, }) diff --git a/internal/handler/hitl_logs.go b/internal/handler/hitl_logs.go index 4d08c926..973c6626 100644 --- a/internal/handler/hitl_logs.go +++ b/internal/handler/hitl_logs.go @@ -10,6 +10,7 @@ import ( "time" "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/security" "github.com/gin-gonic/gin" ) @@ -163,6 +164,7 @@ func (h *AgentHandler) ListHITLLogs(c *gin.Context) { q, args := h.buildHitlListQuery(true) q, args = h.appendHitlListFilters(q, args, c) + q, args = appendConversationAccessSQL(q, args, "conversation_id", notificationAccessFromContext(c)) total, err := h.countHitlQuery(q, args) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) @@ -207,6 +209,7 @@ func (h *AgentHandler) DeleteHITLLogs(c *gin.Context) { if request.All { where, args := h.buildHitlLogsWhere(true) where, args = h.appendHitlListFilters(where, args, c) + where, args = appendConversationAccessSQL(where, args, "conversation_id", notificationAccessFromContext(c)) deleted, err = h.db.DeleteHitlInterruptLogsMatching(where, args) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) @@ -222,7 +225,12 @@ func (h *AgentHandler) DeleteHITLLogs(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "审计日志 ID 列表不能为空"}) return } - deleted, err = h.db.DeleteHitlInterruptLogsByIDs(request.IDs) + ids, filterErr := h.filterAllowedHitlInterruptIDs(c, request.IDs) + if filterErr != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": filterErr.Error()}) + return + } + deleted, err = h.db.DeleteHitlInterruptLogsByIDs(ids) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return @@ -259,5 +267,65 @@ func (h *AgentHandler) GetHITLLog(c *gin.Context) { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return } + if !h.hitlConversationAllowed(c, cid) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } c.JSON(http.StatusOK, hitlInterruptRowToMap(rowID, cid, mode, toolName, toolCallID, payload, rowStatus, decidedBy, messageID, decision, comment, createdAt, decidedAt)) } + +func (h *AgentHandler) filterAllowedHitlInterruptIDs(c *gin.Context, ids []string) ([]string, error) { + clean := make([]string, 0, len(ids)) + seen := map[string]struct{}{} + for _, id := range ids { + id = strings.TrimSpace(id) + if id == "" { + continue + } + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + clean = append(clean, id) + } + if len(clean) == 0 { + return clean, nil + } + query := `SELECT id, conversation_id FROM hitl_interrupts WHERE id IN (` + buildPlaceholders(len(clean)) + `)` + args := make([]interface{}, 0, len(clean)) + for _, id := range clean { + args = append(args, id) + } + rows, err := h.db.Query(query, args...) + if err != nil { + return nil, err + } + defer rows.Close() + allowed := make([]string, 0, len(clean)) + for rows.Next() { + var id, conversationID string + if err := rows.Scan(&id, &conversationID); err != nil { + continue + } + if h.hitlConversationAllowed(c, conversationID) { + allowed = append(allowed, id) + } + } + return allowed, rows.Err() +} + +func (h *AgentHandler) hitlInterruptAllowed(c *gin.Context, interruptID string) bool { + var conversationID string + if err := h.db.QueryRow(`SELECT conversation_id FROM hitl_interrupts WHERE id = ?`, strings.TrimSpace(interruptID)).Scan(&conversationID); err != nil { + return false + } + return h.hitlConversationAllowed(c, conversationID) +} + +func (h *AgentHandler) hitlConversationAllowed(c *gin.Context, conversationID string) bool { + session, ok := security.CurrentSession(c) + if !ok { + return false + } + return h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", conversationID) +} diff --git a/internal/handler/notification.go b/internal/handler/notification.go index 8871e944..b83c8e93 100644 --- a/internal/handler/notification.go +++ b/internal/handler/notification.go @@ -9,6 +9,8 @@ import ( "time" "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + "github.com/gin-gonic/gin" "go.uber.org/zap" ) @@ -89,6 +91,10 @@ func normalizedSinceSec(sinceMs int64) int64 { return 0 } +func ptrTime(t time.Time) *time.Time { + return &t +} + func normalizeSinceMs(raw int64) int64 { if raw > 0 { return raw @@ -126,8 +132,82 @@ func i18nText(english bool, zh string, en string) string { return zh } -func (h *NotificationHandler) loadPendingHITLItems(limit int, english bool) ([]NotificationSummaryItem, error) { - rows, err := h.db.Query(` +func notificationAccessFromContext(c *gin.Context) database.RBACListAccess { + session, ok := security.CurrentSession(c) + if !ok { + return database.RBACListAccess{} + } + return database.RBACListAccess{UserID: session.UserID, Scope: session.Scope} +} + +func appendConversationAccessSQL(query string, args []interface{}, column string, access database.RBACListAccess) (string, []interface{}) { + userID := strings.TrimSpace(access.UserID) + if access.Scope == database.RBACScopeAll { + return query, args + } + if userID == "" { + return query + ` AND 1=0`, args + } + query += ` AND ` + column + ` IS NOT NULL AND ` + column + ` <> '' AND ( + EXISTS (SELECT 1 FROM conversations c WHERE c.id = ` + column + ` AND c.owner_user_id = ?) + OR EXISTS ( + SELECT 1 FROM rbac_resource_assignments ra + WHERE ra.user_id = ? AND ra.resource_type = 'conversation' AND ra.resource_id = ` + column + ` + ) + OR EXISTS ( + SELECT 1 FROM conversations c + JOIN projects p ON p.id = c.project_id + WHERE c.id = ` + column + ` AND p.owner_user_id = ? + ) + OR EXISTS ( + SELECT 1 FROM conversations c + JOIN rbac_resource_assignments pra ON pra.resource_id = c.project_id + WHERE c.id = ` + column + ` AND pra.user_id = ? AND pra.resource_type = 'project' + ) + )` + args = append(args, userID, userID, userID, userID) + return query, args +} + +func appendVulnerabilityNotificationAccessSQL(query string, args []interface{}, access database.RBACListAccess) (string, []interface{}) { + userID := strings.TrimSpace(access.UserID) + if access.Scope == database.RBACScopeAll { + return query, args + } + if userID == "" { + return query + ` AND 1=0`, args + } + query += ` AND ( + owner_user_id = ? + OR EXISTS ( + SELECT 1 FROM rbac_resource_assignments ra + WHERE ra.user_id = ? AND ra.resource_type = 'vulnerability' AND ra.resource_id = vulnerabilities.id + ) + OR ( + project_id IS NOT NULL AND project_id <> '' AND ( + EXISTS (SELECT 1 FROM projects p WHERE p.id = vulnerabilities.project_id AND p.owner_user_id = ?) + OR EXISTS ( + SELECT 1 FROM rbac_resource_assignments pra + WHERE pra.user_id = ? AND pra.resource_type = 'project' AND pra.resource_id = vulnerabilities.project_id + ) + ) + ) + OR ( + conversation_id IS NOT NULL AND conversation_id <> '' AND ( + EXISTS (SELECT 1 FROM conversations c WHERE c.id = vulnerabilities.conversation_id AND c.owner_user_id = ?) + OR EXISTS ( + SELECT 1 FROM rbac_resource_assignments cra + WHERE cra.user_id = ? AND cra.resource_type = 'conversation' AND cra.resource_id = vulnerabilities.conversation_id + ) + ) + ) + )` + args = append(args, userID, userID, userID, userID, userID, userID) + return query, args +} + +func (h *NotificationHandler) loadPendingHITLItems(limit int, english bool, access database.RBACListAccess) ([]NotificationSummaryItem, error) { + query := ` SELECT id, conversation_id, @@ -135,9 +215,14 @@ func (h *NotificationHandler) loadPendingHITLItems(limit int, english bool) ([]N COALESCE(CAST(strftime('%s', created_at) AS INTEGER), 0) FROM hitl_interrupts WHERE status = 'pending' - ORDER BY created_at DESC + ` + args := []interface{}{} + query, args = appendConversationAccessSQL(query, args, "conversation_id", access) + query += ` ORDER BY created_at DESC LIMIT ? - `, limit) + ` + args = append(args, limit) + rows, err := h.db.Query(query, args...) if err != nil { return nil, err } @@ -170,9 +255,9 @@ func (h *NotificationHandler) loadPendingHITLItems(limit int, english bool) ([]N return items, nil } -func (h *NotificationHandler) loadVulnerabilityItems(sinceMs int64, limit int, english bool) ([]NotificationSummaryItem, map[string]int, error) { +func (h *NotificationHandler) loadVulnerabilityItems(sinceMs int64, limit int, english bool, access database.RBACListAccess) ([]NotificationSummaryItem, map[string]int, error) { sinceSec := normalizedSinceSec(sinceMs) - rows, err := h.db.Query(` + query := ` SELECT id, title, @@ -181,9 +266,15 @@ func (h *NotificationHandler) loadVulnerabilityItems(sinceMs int64, limit int, e COALESCE(CAST(strftime('%s', created_at) AS INTEGER), 0) FROM vulnerabilities WHERE CAST(strftime('%s', created_at) AS INTEGER) > ? + ` + args := []interface{}{sinceSec} + query, args = appendVulnerabilityNotificationAccessSQL(query, args, access) + query += ` ORDER BY created_at DESC LIMIT ? - `, sinceSec, limit) + ` + args = append(args, limit) + rows, err := h.db.Query(query, args...) if err != nil { return nil, nil, err } @@ -241,29 +332,23 @@ func (h *NotificationHandler) loadVulnerabilityItems(sinceMs int64, limit int, e } // loadC2SessionOnlineEvents 新会话上线（c2_events：session + critical，与 Manager.IngestCheckIn 一致） -func (h *NotificationHandler) loadC2SessionOnlineEvents(sinceMs int64, limit int, english bool) ([]NotificationSummaryItem, int, error) { +func (h *NotificationHandler) loadC2SessionOnlineEvents(sinceMs int64, limit int, english bool, access database.RBACListAccess) ([]NotificationSummaryItem, int, error) { sinceSec := normalizedSinceSec(sinceMs) - rows, err := h.db.Query(` - SELECT id, message, COALESCE(session_id, ''), - COALESCE(CAST(strftime('%s', created_at) AS INTEGER), 0) - FROM c2_events - WHERE category = 'session' AND level = 'critical' - AND CAST(strftime('%s', created_at) AS INTEGER) > ? - ORDER BY created_at DESC - LIMIT ? - `, sinceSec, limit) + events, err := h.db.ListC2EventsForAccess(database.ListC2EventsFilter{ + Category: "session", + Level: "critical", + Since: ptrTime(time.Unix(sinceSec, 0)), + Limit: limit, + }, access) if err != nil { return nil, 0, err } - defer rows.Close() items := make([]NotificationSummaryItem, 0, limit) - for rows.Next() { - var id, message, sessionID string - var createdSec int64 - if err := rows.Scan(&id, &message, &sessionID, &createdSec); err != nil { + for _, e := range events { + if e == nil { continue } - desc := strings.TrimSpace(message) + desc := strings.TrimSpace(e.Message) if len(desc) > 220 { desc = desc[:200] + "…" } @@ -271,19 +356,19 @@ func (h *NotificationHandler) loadC2SessionOnlineEvents(sinceMs int64, limit int desc = i18nText(english, "新会话已建立", "A new session was created") } items = append(items, NotificationSummaryItem{ - ID: "c2evt:" + id, + ID: "c2evt:" + e.ID, Level: "p0", Type: "c2_session_online", Title: i18nText(english, "C2 新会话上线", "C2 new session online"), Desc: desc, - Ts: unixSecToRFC3339(createdSec), + Ts: e.CreatedAt.UTC().Format(time.RFC3339), Count: 1, Actionable: false, Read: false, - SessionID: sessionID, + SessionID: e.SessionID, }) } - return items, len(items), rows.Err() + return items, len(items), nil } func (h *NotificationHandler) loadFailedExecutionItems(sinceMs int64, limit int, english bool) ([]NotificationSummaryItem, int, error) { @@ -331,7 +416,7 @@ func (h *NotificationHandler) loadFailedExecutionItems(sinceMs int64, limit int, return items, count, nil } -func (h *NotificationHandler) summarizeLongRunningTasks(threshold time.Duration, english bool) ([]NotificationSummaryItem, int) { +func (h *NotificationHandler) summarizeLongRunningTasks(threshold time.Duration, english bool, access database.RBACListAccess) ([]NotificationSummaryItem, int) { if h.agentHandler == nil || h.agentHandler.tasks == nil { return nil, 0 } @@ -342,6 +427,9 @@ func (h *NotificationHandler) summarizeLongRunningTasks(threshold time.Duration, if t == nil { continue } + if !h.notificationConversationAllowed(access, t.ConversationID) { + continue + } if now.Sub(t.StartedAt) >= threshold { items = append(items, NotificationSummaryItem{ ID: "task_long:" + t.ConversationID, @@ -360,7 +448,7 @@ func (h *NotificationHandler) summarizeLongRunningTasks(threshold time.Duration, return items, len(items) } -func (h *NotificationHandler) summarizeCompletedTasksSince(sinceMs int64, limit int, english bool) ([]NotificationSummaryItem, int) { +func (h *NotificationHandler) summarizeCompletedTasksSince(sinceMs int64, limit int, english bool, access database.RBACListAccess) ([]NotificationSummaryItem, int) { if h.agentHandler == nil || h.agentHandler.tasks == nil { return nil, 0 } @@ -371,6 +459,9 @@ func (h *NotificationHandler) summarizeCompletedTasksSince(sinceMs int64, limit if t == nil { continue } + if !h.notificationConversationAllowed(access, t.ConversationID) { + continue + } if t.CompletedAt.After(since) { items = append(items, NotificationSummaryItem{ ID: "task_completed:" + t.ConversationID + ":" + strconv.FormatInt(t.CompletedAt.Unix(), 10), @@ -626,30 +717,57 @@ func (h *NotificationHandler) GetSummary(c *gin.Context) { if limit > 200 { limit = 200 } + access := notificationAccessFromContext(c) - hitlItems, err := h.loadPendingHITLItems(limit, english) - if err != nil { - h.logger.Warn("加载 HITL 通知失败", zap.Error(err)) - c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to summarize hitl notifications"}) - return + hitlItems := []NotificationSummaryItem{} + if security.SessionHasPermission(c, "hitl:read") { + var err error + hitlItems, err = h.loadPendingHITLItems(limit, english, access) + if err != nil { + h.logger.Warn("加载 HITL 通知失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to summarize hitl notifications"}) + return + } } - vulnItems, vulnCounts, err := h.loadVulnerabilityItems(sinceMs, limit, english) - if err != nil { - h.logger.Warn("加载漏洞通知失败", zap.Error(err)) - c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to summarize vulnerabilities"}) - return + vulnItems := []NotificationSummaryItem{} + vulnCounts := map[string]int{ + "newCriticalVulns": 0, + "newHighVulns": 0, + "newMediumVulns": 0, + "newLowVulns": 0, + "newInfoVulns": 0, + } + if security.SessionHasPermission(c, "vulnerability:read") { + var err error + vulnItems, vulnCounts, err = h.loadVulnerabilityItems(sinceMs, limit, english, access) + if err != nil { + h.logger.Warn("加载漏洞通知失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to summarize vulnerabilities"}) + return + } } - c2OnlineItems, c2OnlineCount, err := h.loadC2SessionOnlineEvents(sinceMs, limit, english) - if err != nil { - h.logger.Warn("加载 C2 会话上线通知失败", zap.Error(err)) - c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to summarize c2 session events"}) - return + c2OnlineItems := []NotificationSummaryItem{} + c2OnlineCount := 0 + if security.SessionHasPermission(c, "c2:read") { + var err error + c2OnlineItems, c2OnlineCount, err = h.loadC2SessionOnlineEvents(sinceMs, limit, english, access) + if err != nil { + h.logger.Warn("加载 C2 会话上线通知失败", zap.Error(err)) + c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to summarize c2 session events"}) + return + } } - longRunningItems, longRunningCount := h.summarizeLongRunningTasks(15*time.Minute, english) - completedItems, completedCount := h.summarizeCompletedTasksSince(sinceMs, limit, english) + longRunningItems := []NotificationSummaryItem{} + completedItems := []NotificationSummaryItem{} + longRunningCount := 0 + completedCount := 0 + if security.SessionHasPermission(c, "tasks:read") || security.SessionHasPermission(c, "chat:read") { + longRunningItems, longRunningCount = h.summarizeLongRunningTasks(15*time.Minute, english, access) + completedItems, completedCount = h.summarizeCompletedTasksSince(sinceMs, limit, english, access) + } items := make([]NotificationSummaryItem, 0, len(hitlItems)+len(vulnItems)+len(c2OnlineItems)+len(longRunningItems)+len(completedItems)) items = append(items, hitlItems...) @@ -658,7 +776,7 @@ func (h *NotificationHandler) GetSummary(c *gin.Context) { items = append(items, longRunningItems...) items = append(items, completedItems...) - items, err = h.applyReadStates(items) + items, err := h.applyReadStates(items) if err != nil { h.logger.Warn("加载通知已读状态失败", zap.Error(err)) c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to load notification read states"}) @@ -697,3 +815,11 @@ func (h *NotificationHandler) GetSummary(c *gin.Context) { Items: items, }) } + +func (h *NotificationHandler) notificationConversationAllowed(access database.RBACListAccess, conversationID string) bool { + conversationID = strings.TrimSpace(conversationID) + if conversationID == "" { + return access.Scope == database.RBACScopeAll + } + return h.db.UserCanAccessResource(access.UserID, access.Scope, "conversation", conversationID) +} diff --git a/internal/handler/project.go b/internal/handler/project.go index fb393562..e19e7e38 100644 --- a/internal/handler/project.go +++ b/internal/handler/project.go @@ -9,6 +9,7 @@ import ( "cyberstrike-ai/internal/attackchain" "cyberstrike-ai/internal/database" "cyberstrike-ai/internal/project" + "cyberstrike-ai/internal/security" "github.com/gin-gonic/gin" "go.uber.org/zap" @@ -70,6 +71,10 @@ func (h *ProjectHandler) CreateProject(c *gin.Context) { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return } + if session, ok := security.CurrentSession(c); ok { + _ = h.db.SetResourceOwner("project", created.ID, session.UserID) + _ = h.db.AssignResourceToUser(session.UserID, "project", created.ID) + } c.JSON(http.StatusOK, created) } @@ -82,7 +87,8 @@ func (h *ProjectHandler) GetDashboardSummary(c *gin.Context) { if limit > 50 { limit = 50 } - summary, err := h.db.GetProjectDashboardSummary(limit) + session, _ := security.CurrentSession(c) + summary, err := h.db.GetProjectDashboardSummaryForAccess(limit, session.UserID, session.Scope) if err != nil { h.logger.Error("获取项目仪表盘摘要失败", zap.Error(err)) c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) @@ -106,7 +112,8 @@ func (h *ProjectHandler) ListProjects(c *gin.Context) { if limit > 500 { limit = 500 } - list, err := h.db.ListProjects(status, search, limit, offset) + session, _ := security.CurrentSession(c) + list, err := h.db.ListProjectsForAccess(status, search, limit, offset, session.UserID, session.Scope) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return @@ -114,7 +121,7 @@ func (h *ProjectHandler) ListProjects(c *gin.Context) { if list == nil { list = []*database.Project{} } - total, err := h.db.CountProjects(status, search) + total, err := h.db.CountProjectsForAccess(status, search, session.UserID, session.Scope) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return diff --git a/internal/handler/rbac.go b/internal/handler/rbac.go new file mode 100644 index 00000000..6dc11d36 --- /dev/null +++ b/internal/handler/rbac.go @@ -0,0 +1,376 @@ +package handler + +import ( + "net/http" + "strconv" + "strings" + + "cyberstrike-ai/internal/audit" + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +type RBACHandler struct { + db *database.DB + logger *zap.Logger + audit *audit.Service + auth *security.AuthManager +} + +func NewRBACHandler(db *database.DB, logger *zap.Logger) *RBACHandler { + return &RBACHandler{db: db, logger: logger} +} + +func (h *RBACHandler) SetAudit(s *audit.Service) { + h.audit = s +} + +func (h *RBACHandler) SetAuthManager(m *security.AuthManager) { + h.auth = m +} + +func (h *RBACHandler) Me(c *gin.Context) { + session, _ := security.CurrentSession(c) + c.JSON(http.StatusOK, gin.H{ + "user": gin.H{ + "id": session.UserID, + "username": session.Username, + "display_name": session.DisplayName, + }, + "roles": session.Roles, + "permissions": permissionKeys(session.Permissions), + "scope": session.Scope, + }) +} + +func (h *RBACHandler) Metadata(c *gin.Context) { + roles, err := h.db.ListRBACRoles() + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + rolePermissions := map[string][]string{} + for _, role := range roles { + keys, _ := h.db.ListRBACRolePermissionKeys(role.ID) + rolePermissions[role.ID] = keys + } + c.JSON(http.StatusOK, gin.H{ + "permissions": security.PermissionCatalog, + "roles": roles, + "role_permissions": rolePermissions, + "scopes": []string{database.RBACScopeAll, database.RBACScopeAssigned, database.RBACScopeOwn}, + }) +} + +func (h *RBACHandler) ListRoles(c *gin.Context) { + roles, err := h.db.ListRBACRoles() + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + out := make([]gin.H, 0, len(roles)) + for _, role := range roles { + keys, _ := h.db.ListRBACRolePermissionKeys(role.ID) + out = append(out, gin.H{ + "id": role.ID, + "name": role.Name, + "description": role.Description, + "scope": role.Scope, + "is_system": role.IsSystem, + "permissions": keys, + "created_at": role.CreatedAt, + "updated_at": role.UpdatedAt, + }) + } + c.JSON(http.StatusOK, gin.H{"roles": out}) +} + +type upsertRBACRoleRequest struct { + ID string `json:"id"` + Name string `json:"name" binding:"required"` + Description string `json:"description"` + Scope string `json:"scope"` + Permissions []string `json:"permissions"` +} + +func (h *RBACHandler) CreateRole(c *gin.Context) { + var req upsertRBACRoleRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + role, err := h.db.UpsertRBACRole("", req.Name, req.Description, req.Scope, req.Permissions) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "rbac", "create_role", "创建平台角色", "role", role.ID, nil) + } + c.JSON(http.StatusOK, gin.H{"role": role}) +} + +func (h *RBACHandler) UpdateRole(c *gin.Context) { + id := strings.TrimSpace(c.Param("id")) + existing, err := h.db.GetRBACRoleByID(id) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "角色不存在"}) + return + } + if existing.IsSystem { + c.JSON(http.StatusBadRequest, gin.H{"error": "系统内置角色不可修改，请创建自定义角色"}) + return + } + var req upsertRBACRoleRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + role, err := h.db.UpsertRBACRole(id, req.Name, req.Description, req.Scope, req.Permissions) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if existing.IsSystem { + role.IsSystem = true + } + if h.audit != nil { + h.audit.RecordOK(c, "rbac", "update_role", "更新平台角色", "role", id, nil) + } + if h.auth != nil { + h.auth.RevokeAllSessions() + } + c.JSON(http.StatusOK, gin.H{"role": role}) +} + +func (h *RBACHandler) DeleteRole(c *gin.Context) { + id := strings.TrimSpace(c.Param("id")) + if err := h.db.DeleteRBACRole(id); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "rbac", "delete_role", "删除平台角色", "role", id, nil) + } + if h.auth != nil { + h.auth.RevokeAllSessions() + } + c.JSON(http.StatusOK, gin.H{"success": true}) +} + +func (h *RBACHandler) ListUsers(c *gin.Context) { + users, err := h.db.ListRBACUsers() + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + out := make([]gin.H, 0, len(users)) + for _, user := range users { + roleIDs, _ := h.db.ListRBACUserRoleIDs(user.ID) + out = append(out, gin.H{ + "id": user.ID, + "username": user.Username, + "display_name": user.DisplayName, + "enabled": user.Enabled, + "is_builtin": user.IsBuiltin, + "roles": roleIDs, + "created_at": user.CreatedAt, + "updated_at": user.UpdatedAt, + }) + } + c.JSON(http.StatusOK, gin.H{"users": out}) +} + +type createRBACUserRequest struct { + Username string `json:"username" binding:"required"` + DisplayName string `json:"display_name"` + Password string `json:"password" binding:"required"` + Enabled *bool `json:"enabled"` + Roles []string `json:"roles"` +} + +func (h *RBACHandler) CreateUser(c *gin.Context) { + var req createRBACUserRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if len(strings.TrimSpace(req.Password)) < 8 { + c.JSON(http.StatusBadRequest, gin.H{"error": "密码长度至少需要 8 位"}) + return + } + hash, err := security.HashPassword(req.Password) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + enabled := true + if req.Enabled != nil { + enabled = *req.Enabled + } + user, err := h.db.CreateRBACUser(req.Username, req.DisplayName, hash, enabled, req.Roles) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "rbac", "create_user", "创建平台用户", "user", user.ID, map[string]interface{}{"username": user.Username}) + } + c.JSON(http.StatusOK, gin.H{"user": user}) +} + +type updateRBACUserRequest struct { + DisplayName *string `json:"display_name"` + Password *string `json:"password"` + Enabled *bool `json:"enabled"` + Roles *[]string `json:"roles"` +} + +func (h *RBACHandler) UpdateUser(c *gin.Context) { + id := strings.TrimSpace(c.Param("id")) + user, err := h.db.GetRBACUserByID(id) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "用户不存在"}) + return + } + var req updateRBACUserRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + displayName := user.DisplayName + if req.DisplayName != nil { + displayName = *req.DisplayName + } + if err := h.db.UpdateRBACUser(id, displayName, req.Enabled, req.Roles); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if req.Password != nil && strings.TrimSpace(*req.Password) != "" { + if len(strings.TrimSpace(*req.Password)) < 8 { + c.JSON(http.StatusBadRequest, gin.H{"error": "密码长度至少需要 8 位"}) + return + } + hash, err := security.HashPassword(*req.Password) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + if err := h.db.UpdateRBACUserPassword(id, hash); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + } + if h.audit != nil { + h.audit.RecordOK(c, "rbac", "update_user", "更新平台用户", "user", id, nil) + } + if h.auth != nil { + h.auth.RevokeUserSessions(id) + } + updated, _ := h.db.GetRBACUserByID(id) + c.JSON(http.StatusOK, gin.H{"user": updated}) +} + +func (h *RBACHandler) DeleteUser(c *gin.Context) { + id := strings.TrimSpace(c.Param("id")) + if err := h.db.DeleteRBACUser(id); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "rbac", "delete_user", "删除平台用户", "user", id, nil) + } + if h.auth != nil { + h.auth.RevokeUserSessions(id) + } + c.JSON(http.StatusOK, gin.H{"success": true}) +} + +type assignResourceRequest struct { + UserID string `json:"user_id" binding:"required"` + ResourceType string `json:"resource_type" binding:"required"` + ResourceID string `json:"resource_id"` + ResourceIDs []string `json:"resource_ids"` +} + +func (h *RBACHandler) AssignResource(c *gin.Context) { + var req assignResourceRequest + if err := c.ShouldBindJSON(&req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + resourceIDs := append([]string(nil), req.ResourceIDs...) + if strings.TrimSpace(req.ResourceID) != "" { + resourceIDs = append(resourceIDs, req.ResourceID) + } + if len(resourceIDs) == 0 { + c.JSON(http.StatusBadRequest, gin.H{"error": "至少需要一个资源 ID"}) + return + } + created, err := h.db.AssignResourcesToUser(req.UserID, req.ResourceType, resourceIDs) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + for _, resourceID := range resourceIDs { + h.audit.RecordOK(c, "rbac", "assign_resource", "授权资源访问", req.ResourceType, strings.TrimSpace(resourceID), map[string]interface{}{"user_id": req.UserID}) + } + } + c.JSON(http.StatusOK, gin.H{ + "success": true, + "requested": len(resourceIDs), + "created": created, + "skipped": int64(len(resourceIDs)) - created, + }) +} + +func (h *RBACHandler) ListResourceAssignments(c *gin.Context) { + rows, err := h.db.ListRBACResourceAssignments(c.Query("user_id")) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"assignments": rows}) +} + +func (h *RBACHandler) ListAssignableResources(c *gin.Context) { + limit, _ := strconv.Atoi(c.DefaultQuery("limit", "50")) + if limit <= 0 || limit > 50 { + limit = 50 + } + offset, _ := strconv.Atoi(c.DefaultQuery("offset", "0")) + if offset < 0 { + offset = 0 + } + resources, err := h.db.ListAssignableRBACResourcesPage(c.Query("type"), c.Query("q"), limit+1, offset) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + hasMore := len(resources) > limit + if hasMore { + resources = resources[:limit] + } + c.JSON(http.StatusOK, gin.H{ + "resources": resources, + "has_more": hasMore, + "limit": limit, + "offset": offset, + }) +} + +func (h *RBACHandler) DeleteResourceAssignment(c *gin.Context) { + id := strings.TrimSpace(c.Param("id")) + if err := h.db.DeleteRBACResourceAssignment(id); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if h.audit != nil { + h.audit.RecordOK(c, "rbac", "delete_resource_assignment", "撤销资源授权", "resource_assignment", id, nil) + } + c.JSON(http.StatusOK, gin.H{"success": true}) +} diff --git a/internal/handler/rbac_test.go b/internal/handler/rbac_test.go new file mode 100644 index 00000000..f4986983 --- /dev/null +++ b/internal/handler/rbac_test.go @@ -0,0 +1,118 @@ +package handler + +import ( + "bytes" + "cyberstrike-ai/internal/database" + "encoding/json" + "net/http" + "net/http/httptest" + "path/filepath" + "testing" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +func TestRBACAssignResourceBatchIsAtomicAndLegacyCompatible(t *testing.T) { + gin.SetMode(gin.TestMode) + db, err := database.NewDB(filepath.Join(t.TempDir(), "rbac-handler.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + + user, err := db.CreateRBACUser("api-member", "API Member", "hash", true, nil) + if err != nil { + t.Fatal(err) + } + p1, _ := db.CreateProject(&database.Project{Name: "p1"}) + p2, _ := db.CreateProject(&database.Project{Name: "p2"}) + p3, _ := db.CreateProject(&database.Project{Name: "p3"}) + + h := NewRBACHandler(db, zap.NewNop()) + router := gin.New() + router.POST("/api/rbac/resource-assignments", h.AssignResource) + + batch := performRBACJSONRequest(t, router, map[string]interface{}{ + "user_id": user.ID, "resource_type": "project", "resource_ids": []string{p1.ID, p2.ID}, + }) + if batch.Code != http.StatusOK { + t.Fatalf("batch status = %d, body = %s", batch.Code, batch.Body.String()) + } + var batchBody map[string]interface{} + if err := json.Unmarshal(batch.Body.Bytes(), &batchBody); err != nil { + t.Fatal(err) + } + if batchBody["created"] != float64(2) { + t.Fatalf("batch response = %#v, want created=2", batchBody) + } + + invalid := performRBACJSONRequest(t, router, map[string]interface{}{ + "user_id": user.ID, "resource_type": "project", "resource_ids": []string{p3.ID, "missing"}, + }) + if invalid.Code != http.StatusBadRequest { + t.Fatalf("invalid status = %d, body = %s", invalid.Code, invalid.Body.String()) + } + rows, err := db.ListRBACResourceAssignments(user.ID) + if err != nil { + t.Fatal(err) + } + if len(rows) != 2 { + t.Fatalf("failed batch persisted partial data: %#v", rows) + } + + legacy := performRBACJSONRequest(t, router, map[string]interface{}{ + "user_id": user.ID, "resource_type": "project", "resource_id": p3.ID, + }) + if legacy.Code != http.StatusOK { + t.Fatalf("legacy status = %d, body = %s", legacy.Code, legacy.Body.String()) + } +} + +func TestRBACAssignableResourcesArePaged(t *testing.T) { + gin.SetMode(gin.TestMode) + db, err := database.NewDB(filepath.Join(t.TempDir(), "rbac-picker.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + + for _, name := range []string{"p1", "p2", "p3"} { + if _, err := db.CreateProject(&database.Project{Name: name}); err != nil { + t.Fatal(err) + } + } + h := NewRBACHandler(db, zap.NewNop()) + router := gin.New() + router.GET("/api/rbac/resources", h.ListAssignableResources) + + request := httptest.NewRequest(http.MethodGet, "/api/rbac/resources?type=project&limit=2&offset=0", nil) + recorder := httptest.NewRecorder() + router.ServeHTTP(recorder, request) + if recorder.Code != http.StatusOK { + t.Fatalf("status = %d, body = %s", recorder.Code, recorder.Body.String()) + } + var body struct { + Resources []database.RBACResourceOption `json:"resources"` + HasMore bool `json:"has_more"` + } + if err := json.Unmarshal(recorder.Body.Bytes(), &body); err != nil { + t.Fatal(err) + } + if len(body.Resources) != 2 || !body.HasMore { + t.Fatalf("page = %#v, has_more = %v; want two rows and another page", body.Resources, body.HasMore) + } +} + +func performRBACJSONRequest(t *testing.T, router http.Handler, payload map[string]interface{}) *httptest.ResponseRecorder { + t.Helper() + body, err := json.Marshal(payload) + if err != nil { + t.Fatal(err) + } + req := httptest.NewRequest(http.MethodPost, "/api/rbac/resource-assignments", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + recorder := httptest.NewRecorder() + router.ServeHTTP(recorder, req) + return recorder +} diff --git a/internal/handler/vulnerability.go b/internal/handler/vulnerability.go index a81c88a6..089748a3 100644 --- a/internal/handler/vulnerability.go +++ b/internal/handler/vulnerability.go @@ -9,6 +9,7 @@ import ( "cyberstrike-ai/internal/audit" "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" "github.com/gin-gonic/gin" "go.uber.org/zap" ) @@ -60,6 +61,16 @@ func (h *VulnerabilityHandler) CreateVulnerability(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) return } + if session, ok := security.CurrentSession(c); ok && session.Scope != database.RBACScopeAll { + if strings.TrimSpace(req.ConversationID) != "" && !h.db.UserCanAccessResource(session.UserID, session.Scope, "conversation", strings.TrimSpace(req.ConversationID)) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权在该对话下创建漏洞"}) + return + } + if strings.TrimSpace(req.ProjectID) != "" && !h.db.UserCanAccessResource(session.UserID, session.Scope, "project", strings.TrimSpace(req.ProjectID)) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权在该项目下创建漏洞"}) + return + } + } vuln := &database.Vulnerability{ ConversationID: req.ConversationID, @@ -86,6 +97,10 @@ func (h *VulnerabilityHandler) CreateVulnerability(c *gin.Context) { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return } + if session, ok := security.CurrentSession(c); ok { + _ = h.db.SetResourceOwner("vulnerability", created.ID, session.UserID) + _ = h.db.AssignResourceToUser(session.UserID, "vulnerability", created.ID) + } if h.audit != nil { h.audit.RecordOK(c, "vulnerability", "create", "创建漏洞记录", "vulnerability", created.ID, map[string]interface{}{ @@ -142,6 +157,7 @@ func (h *VulnerabilityHandler) ListVulnerabilities(c *gin.Context) { offsetStr := c.DefaultQuery("offset", "0") pageStr := c.Query("page") filter := parseVulnerabilityListFilter(c) + access := vulnerabilityAccessFromContext(c) limit, _ := strconv.Atoi(limitStr) offset, _ := strconv.Atoi(offsetStr) @@ -163,7 +179,7 @@ func (h *VulnerabilityHandler) ListVulnerabilities(c *gin.Context) { } // 获取总数 - total, err := h.db.CountVulnerabilities(filter) + total, err := h.db.CountVulnerabilitiesForAccess(filter, access) if err != nil { h.logger.Error("获取漏洞总数失败", zap.Error(err)) // 继续执行，使用0作为总数 @@ -171,7 +187,7 @@ func (h *VulnerabilityHandler) ListVulnerabilities(c *gin.Context) { } // 获取漏洞列表 - vulnerabilities, err := h.db.ListVulnerabilities(limit, offset, filter) + vulnerabilities, err := h.db.ListVulnerabilitiesForAccess(limit, offset, filter, access) if err != nil { h.logger.Error("获取漏洞列表失败", zap.Error(err)) c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) @@ -332,8 +348,9 @@ func (h *VulnerabilityHandler) DeleteVulnerability(c *gin.Context) { // BatchDeleteVulnerabilities 按当前筛选条件批量删除漏洞 func (h *VulnerabilityHandler) BatchDeleteVulnerabilities(c *gin.Context) { filter := parseVulnerabilityListFilter(c) + access := vulnerabilityAccessFromContext(c) - total, err := h.db.CountVulnerabilities(filter) + total, err := h.db.CountVulnerabilitiesForAccess(filter, access) if err != nil { h.logger.Error("统计待删除漏洞失败", zap.Error(err)) c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) @@ -344,7 +361,7 @@ func (h *VulnerabilityHandler) BatchDeleteVulnerabilities(c *gin.Context) { return } - deleted, err := h.db.DeleteVulnerabilitiesByFilter(filter) + deleted, err := h.db.DeleteVulnerabilitiesByFilterForAccess(filter, access) if err != nil { h.logger.Error("批量删除漏洞失败", zap.Error(err), zap.Int("count", total)) c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) @@ -364,8 +381,9 @@ func (h *VulnerabilityHandler) BatchDeleteVulnerabilities(c *gin.Context) { // GetVulnerabilityStats 获取漏洞统计 func (h *VulnerabilityHandler) GetVulnerabilityStats(c *gin.Context) { filter := parseVulnerabilityListFilter(c) + access := vulnerabilityAccessFromContext(c) - stats, err := h.db.GetVulnerabilityStats(filter) + stats, err := h.db.GetVulnerabilityStatsForAccess(filter, access) if err != nil { h.logger.Error("获取漏洞统计失败", zap.Error(err)) c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) @@ -377,7 +395,7 @@ func (h *VulnerabilityHandler) GetVulnerabilityStats(c *gin.Context) { // GetVulnerabilityFilterOptions 获取漏洞筛选建议项 func (h *VulnerabilityHandler) GetVulnerabilityFilterOptions(c *gin.Context) { - options, err := h.db.GetVulnerabilityFilterOptions() + options, err := h.db.GetVulnerabilityFilterOptionsForAccess(vulnerabilityAccessFromContext(c)) if err != nil { h.logger.Error("获取漏洞筛选建议失败", zap.Error(err)) c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) @@ -400,8 +418,9 @@ func (h *VulnerabilityHandler) ExportVulnerabilities(c *gin.Context) { } filter := parseVulnerabilityListFilter(c) + access := vulnerabilityAccessFromContext(c) - total, err := h.db.CountVulnerabilities(filter) + total, err := h.db.CountVulnerabilitiesForAccess(filter, access) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return @@ -411,7 +430,7 @@ func (h *VulnerabilityHandler) ExportVulnerabilities(c *gin.Context) { return } - items, err := h.db.ListVulnerabilities(total, 0, filter) + items, err := h.db.ListVulnerabilitiesForAccess(total, 0, filter, access) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return @@ -477,6 +496,14 @@ func (h *VulnerabilityHandler) ExportVulnerabilities(c *gin.Context) { }) } +func vulnerabilityAccessFromContext(c *gin.Context) database.RBACListAccess { + session, ok := security.CurrentSession(c) + if !ok { + return database.RBACListAccess{} + } + return database.RBACListAccess{UserID: session.UserID, Scope: session.Scope} +} + // appendVulnerabilityMarkdown 单条漏洞的 Markdown 片段（与单文件下载字段对齐，缺省字段不写） func appendVulnerabilityMarkdown(b *strings.Builder, v *database.Vulnerability, titleHeading string) { b.WriteString(fmt.Sprintf("%s %s\n\n", titleHeading, v.Title)) diff --git a/internal/handler/webshell.go b/internal/handler/webshell.go index 87e5b5b1..86113dcc 100644 --- a/internal/handler/webshell.go +++ b/internal/handler/webshell.go @@ -15,6 +15,7 @@ import ( "cyberstrike-ai/internal/audit" "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" "github.com/gin-gonic/gin" "github.com/google/uuid" @@ -144,13 +145,13 @@ func normalizeWindowsCmdPath(p string) string { return strings.ReplaceAll(s, "/", "\\") } -// quotePsSingle 把字符串按 PowerShell 单引号字符串规则转义（内部 ' → ''）。 +// quotePsSingle 把字符串按 PowerShell 单引号字符串规则转义（内部 ' → ”）。 // 供 PowerShell 脚本参数使用，全脚本只用单引号，外层 cmd 再用双引号包裹即可安全传递。 func quotePsSingle(s string) string { return "'" + strings.ReplaceAll(s, "'", "''") + "'" } -// quoteShellSinglePosix 把路径按 POSIX sh 单引号规则转义（内部 ' → '\''） +// quoteShellSinglePosix 把路径按 POSIX sh 单引号规则转义（内部 ' → '\”） func quoteShellSinglePosix(p string) string { if p == "" { return "." @@ -379,7 +380,8 @@ func (h *WebShellHandler) ListConnections(c *gin.Context) { c.JSON(http.StatusServiceUnavailable, gin.H{"error": "database not available"}) return } - list, err := h.db.ListWebshellConnections() + session, _ := security.CurrentSession(c) + list, err := h.db.ListWebshellConnectionsForAccess(session.UserID, session.Scope) if err != nil { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return @@ -434,6 +436,10 @@ func (h *WebShellHandler) CreateConnection(c *gin.Context) { c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) return } + if session, ok := security.CurrentSession(c); ok { + _ = h.db.SetResourceOwner("webshell", conn.ID, session.UserID) + _ = h.db.AssignResourceToUser(session.UserID, "webshell", conn.ID) + } if h.audit != nil { host := req.URL if u, err := url.Parse(req.URL); err == nil { @@ -659,14 +665,15 @@ func (h *WebShellHandler) ListAIConversations(c *gin.Context) { // ExecRequest 执行命令请求（前端传入连接信息 + 命令） type ExecRequest struct { - URL string `json:"url" binding:"required"` - Password string `json:"password"` - Type string `json:"type"` // php, asp, aspx, jsp, custom - Method string `json:"method"` // GET 或 POST，空则默认 POST - CmdParam string `json:"cmd_param"` // 命令参数名，如 cmd/xxx，空则默认 cmd - Encoding string `json:"encoding"` // 响应编码：auto / utf-8 / gbk / gb18030，空则 auto - OS string `json:"os"` // 目标操作系统：auto / linux / windows，当前 exec 不用它，保留字段便于未来扩展 - Command string `json:"command" binding:"required"` + URL string `json:"url" binding:"required"` + Password string `json:"password"` + Type string `json:"type"` // php, asp, aspx, jsp, custom + Method string `json:"method"` // GET 或 POST，空则默认 POST + CmdParam string `json:"cmd_param"` // 命令参数名，如 cmd/xxx，空则默认 cmd + Encoding string `json:"encoding"` // 响应编码：auto / utf-8 / gbk / gb18030，空则 auto + OS string `json:"os"` // 目标操作系统：auto / linux / windows，当前 exec 不用它，保留字段便于未来扩展 + ConnectionID string `json:"connection_id,omitempty"` + Command string `json:"command" binding:"required"` } // ExecResponse 执行命令响应 @@ -714,6 +721,10 @@ func (h *WebShellHandler) Exec(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "url and command are required"}) return } + if !h.webshellConnectionRequestAllowed(c, req.ConnectionID, req.URL) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } parsed, err := url.Parse(req.URL) if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") { @@ -807,6 +818,10 @@ func (h *WebShellHandler) FileOp(c *gin.Context) { c.JSON(http.StatusBadRequest, gin.H{"error": "url and action are required"}) return } + if !h.webshellConnectionRequestAllowed(c, req.ConnectionID, req.URL) { + c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } parsed, err := url.Parse(req.URL) if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") { @@ -883,6 +898,25 @@ func (h *WebShellHandler) FileOp(c *gin.Context) { }) } +func (h *WebShellHandler) webshellConnectionRequestAllowed(c *gin.Context, connectionID, requestURL string) bool { + connectionID = strings.TrimSpace(connectionID) + if connectionID == "" { + return true + } + if h.db == nil { + return false + } + session, ok := security.CurrentSession(c) + if !ok || !h.db.UserCanAccessResource(session.UserID, session.Scope, "webshell", connectionID) { + return false + } + conn, err := h.db.GetWebshellConnection(connectionID) + if err != nil || conn == nil { + return false + } + return strings.TrimSpace(conn.URL) == strings.TrimSpace(requestURL) +} + // ExecWithConnection 在指定 WebShell 连接上执行命令（供 MCP/Agent 等非 HTTP 调用） func (h *WebShellHandler) ExecWithConnection(conn *database.WebShellConnection, command string) (output string, ok bool, errMsg string) { if conn == nil { diff --git a/internal/handler/webshell_rbac_test.go b/internal/handler/webshell_rbac_test.go new file mode 100644 index 00000000..68d17e63 --- /dev/null +++ b/internal/handler/webshell_rbac_test.go @@ -0,0 +1,112 @@ +package handler + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "path/filepath" + "testing" + "time" + + "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/security" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +func TestWebshellExecRequiresConnectionAccessWhenConnectionIDProvided(t *testing.T) { + gin.SetMode(gin.TestMode) + db, user, allowed, hidden := setupWebshellRBACTest(t) + handler := NewWebShellHandler(zap.NewNop(), db) + + w := performWebshellJSON(user, http.MethodPost, "/api/webshell/exec", map[string]interface{}{ + "url": hidden.URL, + "connection_id": hidden.ID, + "command": "id", + }, handler.Exec) + if w.Code != http.StatusForbidden { + t.Fatalf("hidden connection status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String()) + } + + w = performWebshellJSON(user, http.MethodPost, "/api/webshell/exec", map[string]interface{}{ + "url": hidden.URL, + "connection_id": allowed.ID, + "command": "id", + }, handler.Exec) + if w.Code != http.StatusForbidden { + t.Fatalf("mismatched URL status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String()) + } +} + +func TestWebshellFileOpRequiresConnectionAccessWhenConnectionIDProvided(t *testing.T) { + gin.SetMode(gin.TestMode) + db, user, _, hidden := setupWebshellRBACTest(t) + handler := NewWebShellHandler(zap.NewNop(), db) + + w := performWebshellJSON(user, http.MethodPost, "/api/webshell/file", map[string]interface{}{ + "url": hidden.URL, + "connection_id": hidden.ID, + "action": "list", + "path": ".", + }, handler.FileOp) + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String()) + } +} + +func setupWebshellRBACTest(t *testing.T) (*database.DB, *database.RBACUser, *database.WebShellConnection, *database.WebShellConnection) { + t.Helper() + db, err := database.NewDB(filepath.Join(t.TempDir(), "webshell-rbac.db"), zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + user, err := db.CreateRBACUser("operator1", "Operator One", "hash", true, nil) + if err != nil { + t.Fatalf("CreateRBACUser: %v", err) + } + allowed := &database.WebShellConnection{ + ID: "ws_allowed", + URL: "http://127.0.0.1/allowed.php", + Type: "php", + Method: "post", + CmdParam: "cmd", + CreatedAt: time.Now(), + } + hidden := &database.WebShellConnection{ + ID: "ws_hidden", + URL: "http://127.0.0.1/hidden.php", + Type: "php", + Method: "post", + CmdParam: "cmd", + CreatedAt: time.Now(), + } + if err := db.CreateWebshellConnection(allowed); err != nil { + t.Fatalf("CreateWebshellConnection allowed: %v", err) + } + if err := db.CreateWebshellConnection(hidden); err != nil { + t.Fatalf("CreateWebshellConnection hidden: %v", err) + } + if err := db.AssignResourceToUser(user.ID, "webshell", allowed.ID); err != nil { + t.Fatalf("AssignResourceToUser: %v", err) + } + return db, user, allowed, hidden +} + +func performWebshellJSON(user *database.RBACUser, method, path string, body map[string]interface{}, handler gin.HandlerFunc) *httptest.ResponseRecorder { + payload, _ := json.Marshal(body) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(method, path, bytes.NewReader(payload)) + c.Request.Header.Set("Content-Type", "application/json") + c.Set(security.ContextSessionKey, security.Session{ + UserID: user.ID, + Username: user.Username, + Permissions: map[string]bool{"webshell:write": true}, + Scope: database.RBACScopeAssigned, + }) + handler(c) + return w +}