mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-07-13 23:47:24 +02:00
Add files via upload
This commit is contained in:
@@ -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
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
@@ -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 中定义)
|
||||
// ============================================================================
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
})
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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})
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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))
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user