Add files via upload

This commit is contained in:
公明
2026-07-10 16:46:26 +08:00
committed by GitHub
parent 25a76a8c97
commit 0ff8c58fbd
18 changed files with 1461 additions and 122 deletions
+50 -5
View File
@@ -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 {
@@ -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)
}
})
}
+59 -19
View File
@@ -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
}
+8 -4
View File
@@ -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)
}
+88 -14
View File
@@ -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 中定义)
// ============================================================================
+48 -6
View File
@@ -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,
@@ -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)
}
}
}
+118
View File
@@ -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
}
+35 -1
View File
@@ -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)
}
+19 -3
View File
@@ -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,
})
+69 -1
View File
@@ -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)
}
+173 -47
View File
@@ -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_eventssession + 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)
}
+10 -3
View File
@@ -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
+376
View File
@@ -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})
}
+118
View File
@@ -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
}
+35 -8
View File
@@ -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))
+45 -11
View File
@@ -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 {
+112
View File
@@ -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
}