CyberStrikeAI/internal/handler/vulnerability_retest.go

package handler

import (
	"net/http"
	"strings"

	"cyberstrike-ai/internal/audit"
	"cyberstrike-ai/internal/security"
	"github.com/gin-gonic/gin"
	"go.uber.org/zap"
)

// CreateRetest creates a sibling snapshot; the chat endpoint runs the new turn.
func (h *VulnerabilityHandler) CreateRetest(c *gin.Context) {
	var req struct {
		Prompt string `json:"prompt" binding:"required"`
	}
	if err := c.ShouldBindJSON(&req); err != nil || strings.TrimSpace(req.Prompt) == "" || len(req.Prompt) > 20000 {
		c.JSON(http.StatusBadRequest, gin.H{"error": "请输入复测提示词（最多 20000 字节）"})
		return
	}
	if _, ok := security.CurrentSession(c); !ok {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "请先登录"})
		return
	}
	v, err := h.db.GetVulnerability(c.Param("id"))
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "漏洞不存在"})
		return
	}
	source, err := h.db.GetConversation(v.ConversationID)
	if err != nil {
		c.JSON(http.StatusConflict, gin.H{"error": "原始对话已不存在，无法创建复测分支"})
		return
	}
	if session, ok := security.CurrentSession(c); ok {
		if !security.SessionHasPermission(c, "chat:read") || !security.SessionHasPermission(c, "chat:write") || !security.SessionHasPermission(c, "agent:execute") || !security.SessionHasPermission(c, "vulnerability:write") || !h.db.UserCanAccessResource(session.UserID, session.ScopeFor("vulnerability:write"), "vulnerability", v.ID) || !h.db.UserCanAccessResource(session.UserID, session.ScopeFor("chat:read"), "conversation", source.ID) || (source.ProjectID != "" && (!security.SessionHasPermission(c, "project:read") || !h.db.UserCanAccessResource(session.UserID, session.ScopeFor("project:read"), "project", source.ProjectID))) {
			c.JSON(http.StatusForbidden, gin.H{"error": "无权创建或执行该漏洞的复测分支"})
			return
		}
	}
	if v.ProjectID != "" && v.ProjectID != source.ProjectID {
		c.JSON(http.StatusConflict, gin.H{"error": "漏洞与原始对话所属项目不一致，请先检查项目归属"})
		return
	}
	if h.taskState != nil {
		if running, _ := h.taskState.ConversationTaskRuntimeState(source.ID); running {
			c.JSON(http.StatusConflict, gin.H{"error": "原始对话仍在执行，请等待完成或停止后再复测"})
			return
		}
	}
	meta := audit.ConversationCreateMetaFromGin(c, "vulnerability_retest")
	conv, err := h.db.CreateVulnerabilityRetest(v, meta)
	if err != nil {
		h.logger.Error("创建复测分支失败", zap.Error(err))
		c.JSON(http.StatusInternalServerError, gin.H{"error": "创建复测分支失败"})
		return
	}
	if session, ok := security.CurrentSession(c); ok {
		if err = h.db.SetResourceOwner("conversation", conv.ID, session.UserID); err == nil {
			err = h.db.AssignResourceToUser(session.UserID, "conversation", conv.ID)
		}
		if err != nil {
			_ = h.db.DeleteConversation(conv.ID)
			c.JSON(http.StatusInternalServerError, gin.H{"error": "设置分支权限失败"})
			return
		}
	}
	// Continue with the user's message unchanged apart from placeholder resolution.
	// Branch linkage and tool authorization are maintained by the backend.
	prompt, err := h.db.ResolveConversationPlaceholders(conv.ID, strings.TrimSpace(req.Prompt))
	if err != nil {
		_ = h.db.DeleteConversation(conv.ID)
		c.JSON(http.StatusInternalServerError, gin.H{"error": "解析复测提示词失败"})
		return
	}
	if err = h.db.SetVulnerabilityRetestPrompt(conv.ID, prompt); err != nil {
		_ = h.db.DeleteConversation(conv.ID)
		c.JSON(http.StatusInternalServerError, gin.H{"error": "保存复测提示词失败"})
		return
	}
	if h.audit != nil {
		h.audit.RecordOK(c, "vulnerability", "retest", "创建漏洞复测分支", "vulnerability", v.ID, map[string]interface{}{"conversation_id": conv.ID, "source_conversation_id": source.ID})
	}
	c.JSON(http.StatusOK, gin.H{"conversation": conv, "prompt": prompt, "vulnerability_id": v.ID, "source_conversation_id": source.ID})
}