nodes.go

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 where a specified CVE is detected in node operating system packages" +
			", grouped by cluster and OS image. Checks OS-level vulnerabilities on cluster nodes." +
			" For comprehensive CVE coverage, also use get_clusters_with_orchestrator_cve (K8s components)" +
			" and get_deployments_for_cve (workloads).",
		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
}