clusters.go

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
}