ROX-31479: Add getting nodes for CVE (#19)

Assisted-by: Claude Code

Author: mtodor

internal/toolsets/vulnerability/nodes.go

@@ -0,0 +1,213 @@
+package vulnerability
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"sort"
+	"strings"
+
+	"github.com/google/jsonschema-go/jsonschema"
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+	"github.com/pkg/errors"
+	v1 "github.com/stackrox/rox/generated/api/v1"
+	"github.com/stackrox/stackrox-mcp/internal/client"
+	"github.com/stackrox/stackrox-mcp/internal/client/auth"
+	"github.com/stackrox/stackrox-mcp/internal/logging"
+	"github.com/stackrox/stackrox-mcp/internal/toolsets"
+	"google.golang.org/grpc"
+)
+
+// getNodesForCVEInput defines the input parameters for get_nodes_for_cve tool.
+type getNodesForCVEInput struct {
+	CVEName         string `json:"cveName"`
+	FilterClusterID string `json:"filterClusterId,omitempty"`
+}
+
+func (input *getNodesForCVEInput) validate() error {
+	if input.CVEName == "" {
+		return errors.New("CVE name is required")
+	}
+
+	return nil
+}
+
+// NodeGroupResult contains aggregated node information by cluster and OS.
+type NodeGroupResult struct {
+	ClusterID       string `json:"clusterId"`
+	ClusterName     string `json:"clusterName"`
+	OperatingSystem string `json:"operatingSystem"`
+	Count           int    `json:"count"`
+}
+
+// getNodesForCVEOutput defines the output structure for get_nodes_for_cve tool.
+type getNodesForCVEOutput struct {
+	NodeGroups []NodeGroupResult `json:"nodeGroups"`
+}
+
+// getNodesForCVETool implements the get_nodes_for_cve tool.
+type getNodesForCVETool struct {
+	name   string
+	client *client.Client
+}
+
+// NewGetNodesForCVETool creates a new get_nodes_for_cve tool.
+func NewGetNodesForCVETool(c *client.Client) toolsets.Tool {
+	return &getNodesForCVETool{
+		name:   "get_nodes_for_cve",
+		client: c,
+	}
+}
+
+// IsReadOnly returns true as this tool only reads data.
+func (t *getNodesForCVETool) IsReadOnly() bool {
+	return true
+}
+
+// GetName returns the tool name.
+func (t *getNodesForCVETool) GetName() string {
+	return t.name
+}
+
+// GetTool returns the MCP Tool definition.
+func (t *getNodesForCVETool) GetTool() *mcp.Tool {
+	return &mcp.Tool{
+		Name:        t.name,
+		Description: "Get aggregated node groups affected by a specific CVE, grouped by cluster and operating system image",
+		InputSchema: getNodesForCVEInputSchema(),
+	}
+}
+
+// getNodesForCVEInputSchema returns the JSON schema for input validation.
+func getNodesForCVEInputSchema() *jsonschema.Schema {
+	schema, err := jsonschema.For[getNodesForCVEInput](nil)
+	if err != nil {
+		logging.Fatal("Could not get jsonschema for get_nodes_for_cve input", err)
+
+		return nil
+	}
+
+	// CVE name is required.
+	schema.Required = []string{"cveName"}
+
+	schema.Properties["cveName"].Description = "CVE name to filter nodes (e.g., CVE-2020-26159)"
+	schema.Properties["filterClusterId"].Description = "Optional cluster ID to filter nodes"
+
+	return schema
+}
+
+// RegisterWith registers the get_nodes_for_cve tool handler with the MCP server.
+func (t *getNodesForCVETool) RegisterWith(server *mcp.Server) {
+	mcp.AddTool(server, t.GetTool(), t.handle)
+}
+
+// buildNodeQuery builds query used to search nodes in StackRox Central.
+// We will quote values to have strict match. Without quote: CVE-2025-10, would match CVE-2025-101.
+func buildNodeQuery(input getNodesForCVEInput) string {
+	queryParts := []string{fmt.Sprintf("CVE:%q", input.CVEName)}
+
+	if input.FilterClusterID != "" {
+		queryParts = append(queryParts, fmt.Sprintf("Cluster ID:%q", input.FilterClusterID))
+	}
+
+	return strings.Join(queryParts, "+")
+}
+
+// aggregateNodeGroups consumes entire stream and aggregates nodes by cluster and OS.
+func aggregateNodeGroups(
+	stream grpc.ServerStreamingClient[v1.ExportNodeResponse],
+) ([]NodeGroupResult, error) {
+	// Map key: "clusterId|osImage"
+	// Map value: NodeGroupResult with count and clusterName.
+	groups := make(map[string]*NodeGroupResult)
+
+	for {
+		resp, err := stream.Recv()
+
+		// Stream ended - no more nodes.
+		if errors.Is(err, io.EOF) {
+			break
+		}
+
+		if err != nil {
+			return nil, errors.Wrap(err, "error receiving from stream")
+		}
+
+		node := resp.GetNode()
+		if node == nil {
+			continue
+		}
+
+		// Create unique key for this cluster+OS combination.
+		key := fmt.Sprintf("%s|%s", node.GetClusterId(), node.GetOsImage())
+		if group, exists := groups[key]; exists {
+			group.Count++
+
+			continue
+		}
+
+		groups[key] = &NodeGroupResult{
+			ClusterID:       node.GetClusterId(),
+			ClusterName:     node.GetClusterName(),
+			OperatingSystem: node.GetOsImage(),
+			Count:           1,
+		}
+	}
+
+	result := make([]NodeGroupResult, 0, len(groups))
+	for _, group := range groups {
+		result = append(result, *group)
+	}
+
+	// Sort for consistent ordering (by clusterId, then OS).
+	sort.Slice(result, func(i, j int) bool {
+		if result[i].ClusterID != result[j].ClusterID {
+			return result[i].ClusterID < result[j].ClusterID
+		}
+
+		return result[i].OperatingSystem < result[j].OperatingSystem
+	})
+
+	return result, nil
+}
+
+// handle is the handler for get_nodes_for_cve tool.
+func (t *getNodesForCVETool) handle(
+	ctx context.Context,
+	req *mcp.CallToolRequest,
+	input getNodesForCVEInput,
+) (*mcp.CallToolResult, *getNodesForCVEOutput, error) {
+	err := input.validate()
+	if err != nil {
+		return nil, nil, err
+	}
+
+	conn, err := t.client.ReadyConn(ctx)
+	if err != nil {
+		return nil, nil, errors.Wrap(err, "unable to connect to server")
+	}
+
+	callCtx := auth.WithMCPRequestContext(ctx, req)
+	nodeClient := v1.NewNodeServiceClient(conn)
+
+	query := buildNodeQuery(input)
+	exportReq := &v1.ExportNodeRequest{
+		Query: query,
+	}
+
+	stream, err := nodeClient.ExportNodes(callCtx, exportReq)
+	if err != nil {
+		return nil, nil, client.NewError(err, "ExportNodes")
+	}
+
+	nodeGroups, err := aggregateNodeGroups(stream)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	output := &getNodesForCVEOutput{
+		NodeGroups: nodeGroups,
+	}
+
+	return nil, output, nil
+}