Add get clusters for CVE

Author: mtodor

internal/toolsets/vulnerability/clusters.go

@@ -0,0 +1,159 @@
+package vulnerability
+
+import (
+	"context"
+	"fmt"
+	"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"
+)
+
+// getClustersForCVEInput defines the input parameters for get_clusters_for_cve tool.
+type getClustersForCVEInput struct {
+	CVEName         string `json:"cveName"`
+	FilterClusterID string `json:"filterClusterId,omitempty"`
+}
+
+func (input *getClustersForCVEInput) validate() error {
+	if input.CVEName == "" {
+		return errors.New("CVE name is required")
+	}
+
+	return nil
+}
+
+// ClusterResult contains cluster information.
+type ClusterResult struct {
+	ClusterID   string `json:"clusterId"`
+	ClusterName string `json:"clusterName"`
+}
+
+// getClustersForCVEOutput defines the output structure for get_clusters_for_cve tool.
+type getClustersForCVEOutput struct {
+	Clusters []ClusterResult `json:"clusters"`
+}
+
+// getClustersForCVETool implements the get_clusters_for_cve tool.
+type getClustersForCVETool struct {
+	name   string
+	client *client.Client
+}
+
+// NewGetClustersForCVETool creates a new get_clusters_for_cve tool.
+func NewGetClustersForCVETool(c *client.Client) toolsets.Tool {
+	return &getClustersForCVETool{
+		name:   "get_clusters_for_cve",
+		client: c,
+	}
+}
+
+// IsReadOnly returns true as this tool only reads data.
+func (t *getClustersForCVETool) IsReadOnly() bool {
+	return true
+}
+
+// GetName returns the tool name.
+func (t *getClustersForCVETool) GetName() string {
+	return t.name
+}
+
+// GetTool returns the MCP Tool definition.
+func (t *getClustersForCVETool) GetTool() *mcp.Tool {
+	return &mcp.Tool{
+		Name:        t.name,
+		Description: "Get list of clusters affected by a specific CVE",
+		InputSchema: getClustersForCVEInputSchema(),
+	}
+}
+
+// getClustersForCVEInputSchema returns the JSON schema for input validation.
+func getClustersForCVEInputSchema() *jsonschema.Schema {
+	schema, err := jsonschema.For[getClustersForCVEInput](nil)
+	if err != nil {
+		logging.Fatal("Could not get jsonschema for get_clusters_for_cve input", err)
+
+		return nil
+	}
+
+	// CVE name is required.
+	schema.Required = []string{"cveName"}
+
+	schema.Properties["cveName"].Description = "CVE name to filter clusters (e.g., CVE-2021-44228)"
+	schema.Properties["filterClusterId"].Description = "Optional cluster ID to verify if a specific cluster is affected"
+
+	return schema
+}
+
+// RegisterWith registers the get_clusters_for_cve tool handler with the MCP server.
+func (t *getClustersForCVETool) RegisterWith(server *mcp.Server) {
+	mcp.AddTool(server, t.GetTool(), t.handle)
+}
+
+// buildClusterQuery builds query string for filtering clusters by CVE.
+// We quote values for exact match (CVE-2025-10 won't match CVE-2025-101).
+func buildClusterQuery(input getClustersForCVEInput) 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, "+")
+}
+
+// handle is the handler for get_clusters_for_cve tool.
+func (t *getClustersForCVETool) handle(
+	ctx context.Context,
+	req *mcp.CallToolRequest,
+	input getClustersForCVEInput,
+) (*mcp.CallToolResult, *getClustersForCVEOutput, 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)
+
+	clustersClient := v1.NewClustersServiceClient(conn)
+
+	query := buildClusterQuery(input)
+
+	resp, err := clustersClient.GetClusters(callCtx, &v1.GetClustersRequest{
+		Query: query,
+	})
+	if err != nil {
+		return nil, nil, client.NewError(err, "GetClusters")
+	}
+
+	clusters := make([]ClusterResult, 0, len(resp.GetClusters()))
+	for _, cluster := range resp.GetClusters() {
+		clusters = append(clusters, ClusterResult{
+			ClusterID:   cluster.GetId(),
+			ClusterName: cluster.GetName(),
+		})
+	}
+
+	// Sort by cluster ID for deterministic output.
+	sort.Slice(clusters, func(i, j int) bool {
+		return clusters[i].ClusterID < clusters[j].ClusterID
+	})
+
+	output := &getClustersForCVEOutput{
+		Clusters: clusters,
+	}
+
+	return nil, output, nil
+}