CyberStrikeAI/internal/handler/vulnerability_retest_test.go

package handler

import (
	"bytes"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strings"
	"testing"

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

func TestCreateVulnerabilityRetestPromptAndPermissions(t *testing.T) {
	gin.SetMode(gin.TestMode)
	db, err := database.NewDB(filepath.Join(t.TempDir(), "retest.db"), zap.NewNop())
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	user, err := db.CreateRBACUser("retester", "Retester", "hash", true, nil)
	if err != nil {
		t.Fatal(err)
	}
	source, err := db.CreateConversation("source", database.ConversationCreateMeta{})
	if err != nil {
		t.Fatal(err)
	}
	hitl := NewHITLManager(db, zap.NewNop())
	if err = hitl.EnsureSchema(); err != nil {
		t.Fatal(err)
	}
	if err = hitl.SaveConversationConfig(source.ID, &HITLRequest{Enabled: true, Mode: "approval", Reviewer: "audit_agent", SensitiveTools: []string{"execute"}, TimeoutSeconds: 90}); err != nil {
		t.Fatal(err)
	}
	v, err := db.CreateVulnerability(&database.Vulnerability{ConversationID: source.ID, Title: "finding", Severity: "high", Status: "open"})
	if err != nil {
		t.Fatal(err)
	}
	if err = db.AssignResourceToUser(user.ID, "vulnerability", v.ID); err != nil {
		t.Fatal(err)
	}
	h := NewVulnerabilityHandler(db, zap.NewNop())
	permissions := map[string]bool{"vulnerability:write": true, "chat:read": true, "chat:write": true, "agent:execute": true}
	request := func(prompt string) *httptest.ResponseRecorder {
		body, _ := json.Marshal(map[string]string{"prompt": prompt})
		w := httptest.NewRecorder()
		c, _ := gin.CreateTestContext(w)
		c.Request = httptest.NewRequest(http.MethodPost, "/api/vulnerabilities/"+v.ID+"/retest", bytes.NewReader(body))
		c.Request.Header.Set("Content-Type", "application/json")
		c.Params = gin.Params{{Key: "id", Value: v.ID}}
		c.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: database.RBACScopeAssigned, Permissions: permissions})
		h.CreateRetest(c)
		return w
	}
	w := request("check {{vulnerability_id}} in {{conversation_id}} from {{source_conversation_id}}")
	if w.Code != http.StatusForbidden {
		t.Fatalf("missing source access allowed: %d %s", w.Code, w.Body.String())
	}
	if err = db.AssignResourceToUser(user.ID, "conversation", source.ID); err != nil {
		t.Fatal(err)
	}
	h.SetTaskStateProvider(staticConversationTaskState{running: true})
	if w = request("check"); w.Code != http.StatusConflict {
		t.Fatalf("running source accepted: %d", w.Code)
	}
	h.SetTaskStateProvider(staticConversationTaskState{running: false})
	w = request("check {{vulnerability_id}} in {{conversation_id}} from {{source_conversation_id}}")
	if w.Code != http.StatusOK {
		t.Fatalf("retest failed: %d %s", w.Code, w.Body.String())
	}
	var result struct {
		Conversation database.Conversation `json:"conversation"`
		Prompt       string                `json:"prompt"`
	}
	if err = json.Unmarshal(w.Body.Bytes(), &result); err != nil {
		t.Fatal(err)
	}
	want := "check " + v.ID + " in " + result.Conversation.ID + " from " + source.ID
	if result.Prompt != want || result.Conversation.ID == source.ID || !db.UserCanAccessResource(user.ID, database.RBACScopeAssigned, "conversation", result.Conversation.ID) {
		t.Fatalf("bad response: %+v", result)
	}
	retest, err := db.GetVulnerabilityRetest(result.Conversation.ID)
	if err != nil || retest.PendingPrompt != result.Prompt {
		t.Fatalf("prompt not persisted: %+v %v", retest, err)
	}
	read := httptest.NewRecorder()
	readCtx, _ := gin.CreateTestContext(read)
	readCtx.Request = httptest.NewRequest(http.MethodGet, "/api/conversations/"+result.Conversation.ID, nil)
	readCtx.Params = gin.Params{{Key: "id", Value: result.Conversation.ID}}
	NewConversationHandler(db, zap.NewNop()).GetConversation(readCtx)
	var recovered database.Conversation
	if err = json.Unmarshal(read.Body.Bytes(), &recovered); err != nil || recovered.Retest == nil || recovered.Retest.PendingPrompt != result.Prompt {
		t.Fatalf("cannot recover pending prompt: %s %v", read.Body.String(), err)
	}
	inherited, err := hitl.LoadConversationConfig(result.Conversation.ID)
	if err != nil || !inherited.Enabled || inherited.Reviewer != "audit_agent" || inherited.Mode != "approval" || inherited.TimeoutSeconds != 90 || len(inherited.SensitiveTools) != 1 || inherited.SensitiveTools[0] != "execute" {
		t.Fatalf("approval policy not inherited: %+v %v", inherited, err)
	}
	chatRecorder := httptest.NewRecorder()
	chatCtx, _ := gin.CreateTestContext(chatRecorder)
	chatCtx.Request = httptest.NewRequest(http.MethodPost, "/api/eino-agent", nil)
	chatCtx.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: database.RBACScopeAssigned, Permissions: permissions})
	agentHandler := &AgentHandler{db: db, logger: zap.NewNop()}
	req := &ChatRequest{ConversationID: result.Conversation.ID, Message: "continue {{vulnerability_id}} in {{conversation_id}} from {{source_conversation_id}}"}
	prep, err := agentHandler.prepareMultiAgentSession(req, chatCtx, "test")
	if err != nil || !strings.Contains(prep.FinalMessage, "continue "+v.ID+" in "+result.Conversation.ID+" from "+source.ID) {
		t.Fatalf("follow-up placeholders failed: %+v %v", prep, err)
	}
	retest, err = db.GetVulnerabilityRetest(result.Conversation.ID)
	if err != nil || retest.PendingPrompt != "" {
		t.Fatalf("sent prompt still pending: %+v %v", retest, err)
	}
	defaultPrompt := "请帮我对漏洞 {{vulnerability_id}} 进行验证，确认是否已修复。完成后更新该漏洞状态。"
	w = request(defaultPrompt)
	if w.Code != http.StatusOK {
		t.Fatalf("default retest failed: %d %s", w.Code, w.Body.String())
	}
	if err = json.Unmarshal(w.Body.Bytes(), &result); err != nil {
		t.Fatal(err)
	}
	if result.Prompt != strings.ReplaceAll(defaultPrompt, "{{vulnerability_id}}", v.ID) || strings.Contains(result.Prompt, source.ID) || strings.Contains(result.Prompt, result.Conversation.ID) {
		t.Fatalf("branch context injected into user prompt: %q", result.Prompt)
	}
	delete(permissions, "agent:execute")
	if w = request("check"); w.Code != http.StatusForbidden {
		t.Fatal("execution permission not enforced")
	}
	if w = request(" "); w.Code != http.StatusBadRequest {
		t.Fatal("empty prompt allowed")
	}
}