CyberStrikeAI/internal/app/asset_tools_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 TestAssetToolsCRUDQueryAndPageLimit(t *testing.T) {
	db, err := database.NewDB(filepath.Join(t.TempDir(), "asset-tools.db"), zap.NewNop())
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	user, err := db.CreateRBACUser("asset-agent", "Asset Agent", "hash", true, nil)
	if err != nil {
		t.Fatal(err)
	}
	principal := authctx.NewPrincipal(user.ID, user.Username, database.RBACScopeAssigned, map[string]bool{
		"asset:read": true, "asset:write": true, "asset:delete": true,
	})
	ctx := authctx.WithPrincipal(context.Background(), principal)
	server := mcp.NewServer(zap.NewNop())
	server.SetToolAuthorizer(mcpToolAuthorizer(db))
	registerAssetTools(server, db, zap.NewNop())

	wantTools := map[string]bool{
		builtin.ToolCreateAsset: false, builtin.ToolGetAsset: false, builtin.ToolQueryAssets: false,
		builtin.ToolUpdateAsset: false, builtin.ToolDeleteAsset: false, builtin.ToolCompleteAssetScan: false,
	}
	for _, tool := range server.GetAllTools() {
		if _, ok := wantTools[tool.Name]; ok {
			wantTools[tool.Name] = true
		}
	}
	for name, found := range wantTools {
		if !found {
			t.Fatalf("asset tool not registered: %s", name)
		}
	}

	result, _, err := server.CallTool(ctx, builtin.ToolCreateAsset, map[string]interface{}{
		"ip": "192.0.2.42", "port": 443, "protocol": "https", "title": "Before", "tags": []interface{}{"prod"},
	})
	if err != nil || result == nil || result.IsError {
		t.Fatalf("create asset result=%#v err=%v", result, err)
	}
	assets, total, err := db.ListAssets(20, 0, database.AssetListFilter{}, database.RBACListAccess{UserID: user.ID, Scope: database.RBACScopeAssigned})
	if err != nil || total != 1 || len(assets) != 1 {
		t.Fatalf("saved assets total=%d len=%d err=%v", total, len(assets), err)
	}
	id := assets[0].ID

	result, _, err = server.CallTool(ctx, builtin.ToolUpdateAsset, map[string]interface{}{"id": id, "title": "After"})
	if err != nil || result == nil || result.IsError {
		t.Fatalf("update asset result=%#v err=%v", result, err)
	}
	updated, err := db.GetAsset(id, database.RBACListAccess{UserID: user.ID, Scope: database.RBACScopeAssigned})
	if err != nil || updated.Title != "After" || updated.IP != "192.0.2.42" {
		t.Fatalf("partial update lost fields: %#v err=%v", updated, err)
	}

	result, _, err = server.CallTool(ctx, builtin.ToolQueryAssets, map[string]interface{}{
		"sort_by": "last_scan_at", "sort_order": "asc", "page": 1, "page_size": 1,
	})
	if err != nil || result == nil || result.IsError || !strings.Contains(toolResultText(result), "第 1/1 页") || !strings.Contains(toolResultText(result), "last_scan_at=never") {
		t.Fatalf("query asset result=%#v err=%v", result, err)
	}
	result, _, err = server.CallTool(ctx, builtin.ToolQueryAssets, map[string]interface{}{"page_size": agentAssetPageSizeMax + 1})
	if err != nil || result == nil || !result.IsError {
		t.Fatalf("oversized page was accepted: result=%#v err=%v", result, err)
	}

	conversation, err := db.CreateConversation("asset scan", database.ConversationCreateMeta{})
	if err != nil {
		t.Fatal(err)
	}
	if err := db.AssignResourceToUser(user.ID, "conversation", conversation.ID); err != nil {
		t.Fatal(err)
	}
	if _, err := db.CreateVulnerability(&database.Vulnerability{ConversationID: conversation.ID, Title: "finding", Severity: "high", Target: "192.0.2.42"}); err != nil {
		t.Fatal(err)
	}
	scanCtx := mcp.WithMCPConversationID(ctx, conversation.ID)
	result, _, err = server.CallTool(scanCtx, builtin.ToolCompleteAssetScan, map[string]interface{}{"id": id})
	if err != nil || result == nil || result.IsError {
		t.Fatalf("complete scan result=%#v err=%v", result, err)
	}
	scanned, err := db.GetAsset(id, database.RBACListAccess{UserID: user.ID, Scope: database.RBACScopeAssigned})
	if err != nil || scanned.LastScanAt == nil || scanned.LastScanConversationID != conversation.ID || scanned.VulnerabilityCount != 1 {
		t.Fatalf("scan fields not updated: %#v err=%v", scanned, err)
	}

	result, _, err = server.CallTool(ctx, builtin.ToolDeleteAsset, map[string]interface{}{"id": id})
	if err != nil || result == nil || result.IsError {
		t.Fatalf("delete asset result=%#v err=%v", result, err)
	}
	if _, err := db.GetAsset(id, database.RBACListAccess{Scope: database.RBACScopeAll}); err == nil {
		t.Fatal("asset still exists after delete")
	}
}

func toolResultText(result *mcp.ToolResult) string {
	var b strings.Builder
	if result == nil {
		return ""
	}
	for _, content := range result.Content {
		b.WriteString(content.Text)
	}
	return b.String()
}

func TestAssetReadToolsRespectConversationProjectScope(t *testing.T) {
	db, err := database.NewDB(filepath.Join(t.TempDir(), "asset-project-scope.db"), zap.NewNop())
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()

	projectA, err := db.CreateProject(&database.Project{Name: "Project A"})
	if err != nil {
		t.Fatal(err)
	}
	projectB, err := db.CreateProject(&database.Project{Name: "Project B"})
	if err != nil {
		t.Fatal(err)
	}
	assets := []*database.Asset{
		{ProjectID: projectA.ID, IP: "192.0.2.10", Protocol: "https"},
		{ProjectID: projectB.ID, IP: "192.0.2.20", Protocol: "https"},
		{IP: "192.0.2.30", Protocol: "https"},
	}
	if result, err := db.UpsertAssets(assets, "", true); err != nil || result.Created != len(assets) {
		t.Fatalf("seed assets result=%#v err=%v", result, err)
	}

	bound, err := db.CreateConversation("bound", database.ConversationCreateMeta{ProjectID: projectA.ID})
	if err != nil {
		t.Fatal(err)
	}
	unbound, err := db.CreateConversation("unbound", database.ConversationCreateMeta{})
	if err != nil {
		t.Fatal(err)
	}
	principal := authctx.NewPrincipal("admin", "admin", database.RBACScopeAll, map[string]bool{"asset:read": true})
	ctx := authctx.WithPrincipal(context.Background(), principal)
	server := mcp.NewServer(zap.NewNop())
	server.SetToolAuthorizer(mcpToolAuthorizer(db))
	registerAssetTools(server, db, zap.NewNop())

	boundCtx := mcp.WithMCPConversationID(ctx, bound.ID)
	result, _, err := server.CallTool(boundCtx, builtin.ToolQueryAssets, map[string]interface{}{})
	text := toolResultText(result)
	if err != nil || result == nil || result.IsError || !strings.Contains(text, assets[0].ID) || strings.Contains(text, assets[1].ID) || strings.Contains(text, assets[2].ID) {
		t.Fatalf("bound query escaped project scope: result=%#v text=%q err=%v", result, text, err)
	}

	// Even an explicit foreign project_id cannot override the conversation boundary.
	result, _, err = server.CallTool(boundCtx, builtin.ToolQueryAssets, map[string]interface{}{"project_id": projectB.ID})
	text = toolResultText(result)
	if err != nil || result == nil || result.IsError || !strings.Contains(text, assets[0].ID) || strings.Contains(text, assets[1].ID) {
		t.Fatalf("project_id overrode conversation scope: result=%#v text=%q err=%v", result, text, err)
	}

	result, _, err = server.CallTool(boundCtx, builtin.ToolGetAsset, map[string]interface{}{"id": assets[1].ID})
	if err != nil || result == nil || !result.IsError {
		t.Fatalf("bound get read a foreign-project asset: result=%#v err=%v", result, err)
	}

	unboundCtx := mcp.WithMCPConversationID(ctx, unbound.ID)
	result, _, err = server.CallTool(unboundCtx, builtin.ToolQueryAssets, map[string]interface{}{"page_size": 10})
	text = toolResultText(result)
	if err != nil || result == nil || result.IsError || !strings.Contains(text, assets[0].ID) || !strings.Contains(text, assets[1].ID) || !strings.Contains(text, assets[2].ID) {
		t.Fatalf("unbound query did not retain all-assets behavior: result=%#v text=%q err=%v", result, text, err)
	}
}