CyberStrikeAI/internal/app/vulnerability_retest_test.go

package app

import (
	"context"
	"path/filepath"
	"strings"
	"testing"

	"cyberstrike-ai/internal/authctx"
	"cyberstrike-ai/internal/database"
	"cyberstrike-ai/internal/mcp"
	"cyberstrike-ai/internal/mcp/builtin"
	"go.uber.org/zap"
)

func TestVulnerabilityRetestMCPUpdatesOnlyAssignedFinding(t *testing.T) {
	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)
	}
	other, err := db.CreateConversation("other", database.ConversationCreateMeta{})
	if err != nil {
		t.Fatal(err)
	}
	var findings []*database.Vulnerability
	for _, cid := range []string{source.ID, source.ID, other.ID} {
		v, err := db.CreateVulnerability(&database.Vulnerability{ConversationID: cid, Title: "finding", Severity: "high", Status: "open", RetestNotes: "original procedure"})
		if err != nil {
			t.Fatal(err)
		}
		findings = append(findings, v)
		if err = db.AssignResourceToUser(user.ID, "vulnerability", v.ID); err != nil {
			t.Fatal(err)
		}
	}
	branch, err := db.CreateVulnerabilityRetest(findings[0], database.ConversationCreateMeta{})
	if err != nil {
		t.Fatal(err)
	}
	if err = db.AssignResourceToUser(user.ID, "conversation", branch.ID); err != nil {
		t.Fatal(err)
	}
	server := mcp.NewServer(zap.NewNop())
	server.SetToolAuthorizer(mcpToolAuthorizer(db))
	registerVulnerabilityTools(server, db, zap.NewNop())
	principal := authctx.NewPrincipal(user.ID, user.Username, database.RBACScopeAssigned, map[string]bool{"vulnerability:read": true, "vulnerability:write": true})
	ctx := authctx.WithPrincipal(mcp.WithMCPConversationID(context.Background(), branch.ID), principal)
	call := func(id, status, notes string) bool {
		t.Helper()
		result, _, err := server.CallTool(ctx, builtin.ToolUpdateVulnerabilityStatus, map[string]interface{}{"id": id, "status": status, "retest_notes": notes})
		return err == nil && result != nil && !result.IsError
	}
	if call(findings[0].ID, "fixed", "") || call(findings[0].ID, "invalid", "evidence") || call(findings[1].ID, "fixed", "evidence") || call(findings[2].ID, "fixed", "evidence") {
		t.Fatal("invalid or unrelated update succeeded")
	}
	if !call(findings[0].ID, "confirmed", "still reproducible") || !call(findings[0].ID, "fixed", "now access denied with control request succeeding") {
		t.Fatal("assigned finding update denied")
	}
	got, err := db.GetVulnerability(findings[0].ID)
	if err != nil {
		t.Fatal(err)
	}
	if got.Status != "fixed" || got.ConversationID != source.ID || got.RetestNotes != "original procedure" {
		t.Fatalf("bad result: %+v", got)
	}
	if !strings.Contains(got.RetestLog, "still reproducible") || !strings.Contains(got.RetestLog, "now access denied") || !strings.Contains(got.RetestLog, branch.ID) {
		t.Fatalf("retest evidence missing from retest_log: %q", got.RetestLog)
	}
	for _, v := range findings[1:] {
		got, err := db.GetVulnerability(v.ID)
		if err != nil || got.Status != "open" {
			t.Fatalf("unrelated finding changed: %+v %v", got, err)
		}
	}
	result, _, err := server.CallTool(ctx, builtin.ToolGetVulnerability, map[string]interface{}{"id": findings[0].ID})
	if err != nil || result.IsError {
		t.Fatalf("branch cannot read source finding: %v %+v", err, result)
	}
	movedProject, err := db.CreateProject(&database.Project{Name: "different project"})
	if err != nil {
		t.Fatal(err)
	}
	overrideCtx := mcp.WithMCPProjectID(ctx, movedProject.ID)
	if err = mcpToolAuthorizer(db)(overrideCtx, builtin.ToolUpdateVulnerabilityStatus, map[string]interface{}{"id": findings[0].ID}); err == nil {
		t.Fatal("explicit foreign project scope bypassed")
	}
	if err = db.SetConversationProjectID(branch.ID, movedProject.ID); err != nil {
		t.Fatal(err)
	}
	if call(findings[0].ID, "confirmed", "after move") {
		t.Fatal("moving branch bypassed project isolation")
	}
	readonly := authctx.NewPrincipal(user.ID, user.Username, database.RBACScopeAssigned, map[string]bool{"vulnerability:read": true})
	denied := authctx.WithPrincipal(mcp.WithMCPConversationID(context.Background(), branch.ID), readonly)
	if err = mcpToolAuthorizer(db)(denied, builtin.ToolUpdateVulnerabilityStatus, map[string]interface{}{"id": findings[0].ID}); err == nil {
		t.Fatal("read-only principal allowed mutation")
	}
}