ROX-31478: Add tool to get CVE related deployments (#14)

Assisted-by: Claude Code

Author: mtodor

.gitignore

@@ -5,6 +5,7 @@
 
 # Claude Code
 .claude/
+/.mcp.json
 
 # Test output
 /*.out

cmd/stackrox-mcp/main.go

@@ -22,7 +22,7 @@ import (
 func getToolsets(cfg *config.Config, c *client.Client) []toolsets.Toolset {
 	return []toolsets.Toolset{
 		toolsetConfig.NewToolset(cfg, c),
-		toolsetVulnerability.NewToolset(cfg),
+		toolsetVulnerability.NewToolset(cfg, c),
 	}
 }
 

internal/toolsets/vulnerability/tools.go

@@ -2,63 +2,209 @@ package vulnerability
 
 import (
 	"context"
-	"errors"
+	"fmt"
+	"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"
 )
 
-// listClusterCVEsInput defines the input parameters for list_cluster_cves tool.
-type listClusterCVEsInput struct {
-	ClusterID string `json:"clusterId,omitempty"`
+const (
+	defaultLimit = 50
+	maximumLimit = 200.0
+)
+
+type filterPlatformType string
+
+const (
+	filterPlatformNoFilter     filterPlatformType = "NO_FILTER"
+	filterPlatformUserWorkload filterPlatformType = "USER_WORKLOAD"
+	filterPlatformPlatform     filterPlatformType = "PLATFORM"
+)
+
+// getDeploymentsForCVEInput defines the input parameters for get_deployments_for_cve tool.
+type getDeploymentsForCVEInput struct {
+	CVEName         string             `json:"cveName"`
+	FilterClusterID string             `json:"filterClusterId,omitempty"`
+	FilterNamespace string             `json:"filterNamespace,omitempty"`
+	FilterPlatform  filterPlatformType `json:"filterPlatform,omitempty"`
+	Offset          int32              `json:"offset,omitempty"`
+	Limit           int32              `json:"limit,omitempty"`
+}
+
+func (input *getDeploymentsForCVEInput) validate() error {
+	if input.CVEName == "" {
+		return errors.New("CVE name is required")
+	}
+
+	return nil
+}
+
+// DeploymentResult contains deployment information.
+type DeploymentResult struct {
+	Name        string `json:"name"`
+	Namespace   string `json:"namespace"`
+	ClusterID   string `json:"clusterId"`
+	ClusterName string `json:"clusterName"`
 }
 
-// listClusterCVEsOutput defines the output structure for list_cluster_cves tool.
-type listClusterCVEsOutput struct {
-	CVEs []string `json:"cves"`
+// getDeploymentsForCVEOutput defines the output structure for get_deployments_for_cve tool.
+type getDeploymentsForCVEOutput struct {
+	Deployments []DeploymentResult `json:"deployments"`
 }
 
-// listClusterCVEsTool implements the list_cluster_cves tool.
-type listClusterCVEsTool struct {
-	name string
+// getDeploymentsForCVETool implements the get_deployments_for_cve tool.
+type getDeploymentsForCVETool struct {
+	name   string
+	client *client.Client
 }
 
-// NewListClusterCVEsTool creates a new list_cluster_cves tool.
-func NewListClusterCVEsTool() toolsets.Tool {
-	return &listClusterCVEsTool{
-		name: "list_cluster_cves",
+// NewGetDeploymentsForCVETool creates a new get_deployments_for_cve tool.
+func NewGetDeploymentsForCVETool(c *client.Client) toolsets.Tool {
+	return &getDeploymentsForCVETool{
+		name:   "get_deployments_for_cve",
+		client: c,
 	}
 }
 
 // IsReadOnly returns true as this tool only reads data.
-func (t *listClusterCVEsTool) IsReadOnly() bool {
+func (t *getDeploymentsForCVETool) IsReadOnly() bool {
 	return true
 }
 
 // GetName returns the tool name.
-func (t *listClusterCVEsTool) GetName() string {
+func (t *getDeploymentsForCVETool) GetName() string {
 	return t.name
 }
 
 // GetTool returns the MCP Tool definition.
-func (t *listClusterCVEsTool) GetTool() *mcp.Tool {
+func (t *getDeploymentsForCVETool) GetTool() *mcp.Tool {
 	return &mcp.Tool{
-		Name: t.name,
-		//nolint:lll
-		Description: "List CVEs affecting a specific cluster or all clusters in StackRox Central with CVE names, scores, affected images, and deployments",
+		Name:        t.name,
+		Description: "Get list of deployments affected by a specific CVE",
+		InputSchema: getDeploymentsForCVEInputSchema(),
 	}
 }
 
-// RegisterWith registers the list_cluster_cves tool handler with the MCP server.
-func (t *listClusterCVEsTool) RegisterWith(server *mcp.Server) {
+// getDeploymentsForCVEInputSchema returns the JSON schema for input validation.
+func getDeploymentsForCVEInputSchema() *jsonschema.Schema {
+	schema, err := jsonschema.For[getDeploymentsForCVEInput](nil)
+	if err != nil {
+		logging.Fatal("Could not get jsonschema for get_deployments_for_cve input", err)
+
+		return nil
+	}
+
+	// CVE name is required.
+	schema.Required = []string{"cveName"}
+
+	schema.Properties["cveName"].Description = "CVE name to filter deployments (e.g., CVE-2021-44228)"
+	schema.Properties["filterClusterId"].Description = "Optional cluster ID to filter deployments"
+	schema.Properties["filterNamespace"].Description = "Optional namespace to filter deployments"
+
+	schema.Properties["filterPlatform"].Description =
+		fmt.Sprintf("Optional platform filter: %s=no filter, %s=user workload deployments, %s=platform deployments",
+			filterPlatformNoFilter, filterPlatformUserWorkload, filterPlatformPlatform)
+	schema.Properties["filterPlatform"].Default = toolsets.MustJSONMarshal(filterPlatformNoFilter)
+	schema.Properties["filterPlatform"].Enum = []any{
+		filterPlatformNoFilter,
+		filterPlatformUserWorkload,
+		filterPlatformPlatform,
+	}
+
+	schema.Properties["offset"].Description = "Pagination offset (default: 0)"
+	schema.Properties["offset"].Default = toolsets.MustJSONMarshal(0)
+	schema.Properties["limit"].Minimum = jsonschema.Ptr(0.0)
+
+	schema.Properties["limit"].Description = "Pagination limit: minimum: 1, maximum: 200 (default: 50)"
+	schema.Properties["limit"].Default = toolsets.MustJSONMarshal(defaultLimit)
+	schema.Properties["limit"].Minimum = jsonschema.Ptr(1.0)
+	schema.Properties["limit"].Maximum = jsonschema.Ptr(maximumLimit)
+
+	return schema
+}
+
+// RegisterWith registers the get_deployments_for_cve tool handler with the MCP server.
+func (t *getDeploymentsForCVETool) RegisterWith(server *mcp.Server) {
 	mcp.AddTool(server, t.GetTool(), t.handle)
 }
 
-// handle is the placeholder handler for list_cluster_cves tool.
-func (t *listClusterCVEsTool) handle(
-	_ context.Context,
-	_ *mcp.CallToolRequest,
-	_ listClusterCVEsInput,
-) (*mcp.CallToolResult, *listClusterCVEsOutput, error) {
-	return nil, nil, errors.New("list_cluster_cves tool is not yet implemented")
+// buildQuery builds query used to search deployments in StackRox Central.
+// We will quote values to have strict match. Without quote: CVE-2025-10, would match CVE-2025-101.
+func buildQuery(input getDeploymentsForCVEInput) string {
+	queryParts := []string{fmt.Sprintf("CVE:%q", input.CVEName)}
+
+	if input.FilterClusterID != "" {
+		queryParts = append(queryParts, fmt.Sprintf("Cluster ID:%q", input.FilterClusterID))
+	}
+
+	if input.FilterNamespace != "" {
+		queryParts = append(queryParts, fmt.Sprintf("Namespace:%q", input.FilterNamespace))
+	}
+
+	// Add platform filter if provided.
+	switch input.FilterPlatform {
+	case filterPlatformUserWorkload:
+		queryParts = append(queryParts, "Platform Component:0")
+	case filterPlatformPlatform:
+		queryParts = append(queryParts, "Platform Component:1")
+	case filterPlatformNoFilter:
+	}
+
+	return strings.Join(queryParts, "+")
+}
+
+// handle is the handler for get_deployments_for_cve tool.
+func (t *getDeploymentsForCVETool) handle(
+	ctx context.Context,
+	req *mcp.CallToolRequest,
+	input getDeploymentsForCVEInput,
+) (*mcp.CallToolResult, *getDeploymentsForCVEOutput, 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)
+	deploymentClient := v1.NewDeploymentServiceClient(conn)
+
+	listReq := &v1.RawQuery{
+		Query: buildQuery(input),
+		Pagination: &v1.Pagination{
+			Offset: input.Offset,
+			Limit:  input.Limit,
+		},
+	}
+
+	resp, err := deploymentClient.ListDeployments(callCtx, listReq)
+	if err != nil {
+		return nil, nil, client.NewError(err, "ListDeployments")
+	}
+
+	deployments := make([]DeploymentResult, 0, len(resp.GetDeployments()))
+	for _, deployment := range resp.GetDeployments() {
+		deployments = append(deployments, DeploymentResult{
+			Name:        deployment.GetName(),
+			Namespace:   deployment.GetNamespace(),
+			ClusterID:   deployment.GetClusterId(),
+			ClusterName: deployment.GetCluster(),
+		})
+	}
+
+	output := &getDeploymentsForCVEOutput{
+		Deployments: deployments,
+	}
+
+	return nil, output, nil
 }