deployments.go

package vulnerability

import (
	"context"
	"fmt"
	"strings"
	"sync"

	"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/cursor"
	"github.com/stackrox/stackrox-mcp/internal/logging"
	"github.com/stackrox/stackrox-mcp/internal/toolsets"
)

const (
	defaultLimit = 100
)

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"`
	IncludeAffectedImages bool               `json:"includeAffectedImages,omitempty"`
	Cursor                string             `json:"cursor,omitempty"`
}

func (input *getDeploymentsForCVEInput) validate() error {
	if input.CVEName == "" {
		return errors.New("CVE name is required")
	}

	return nil
}

// DeploymentResult contains deployment information with optional image data.
type DeploymentResult struct {
	id string

	Name            string   `json:"name"`
	Namespace       string   `json:"namespace"`
	ClusterID       string   `json:"clusterId"`
	ClusterName     string   `json:"clusterName"`
	AffectedImages  []string `json:"affectedImages,omitempty"`
	ImageFetchError string   `json:"imageFetchError,omitempty"`
}

// getDeploymentsForCVEOutput defines the output structure for get_deployments_for_cve tool.
type getDeploymentsForCVEOutput struct {
	Deployments []DeploymentResult `json:"deployments"`
	NextCursor  string             `json:"nextCursor"`
}

// getDeploymentsForCVETool implements the get_deployments_for_cve tool.
type getDeploymentsForCVETool struct {
	name   string
	client *client.Client
}

// 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 *getDeploymentsForCVETool) IsReadOnly() bool {
	return true
}

// GetName returns the tool name.
func (t *getDeploymentsForCVETool) GetName() string {
	return t.name
}

// GetTool returns the MCP Tool definition.
func (t *getDeploymentsForCVETool) GetTool() *mcp.Tool {
	return &mcp.Tool{
		Name: t.name,
		Description: "Get list of deployments where a specified CVE is detected in application" +
			" or platform container images. Checks user workloads for vulnerabilities." +
			" For complete CVE analysis, also check get_clusters_with_orchestrator_cve (Kubernetes components)" +
			" and get_nodes_for_cve (node OS).",
		InputSchema: getDeploymentsForCVEInputSchema(),
	}
}

// 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["includeAffectedImages"].Description =
		"Whether to include affected image names for each deployment.\n" +
			"WARNING: This may significantly increase response time."
	schema.Properties["includeAffectedImages"].Default = toolsets.MustJSONMarshal(false)

	schema.Properties["cursor"].Description = "Cursor for next page provided by server"

	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)
}

// 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, "+")
}

func getCursor(input *getDeploymentsForCVEInput) (*cursor.Cursor, error) {
	if input.Cursor == "" {
		startCursor, err := cursor.New(0)

		return startCursor, errors.Wrap(err, "error creating starting cursor")
	}

	currCursor, err := cursor.Decode(input.Cursor)
	if err != nil {
		return nil, errors.Wrap(err, "error decoding cursor")
	}

	return currCursor, nil
}

const defaultMaxFetchImageConcurrency = 10

// deploymentEnricher handles parallel enrichment of deployments with image data.
type deploymentEnricher struct {
	imageClient v1.ImageServiceClient
	cveName     string
	semaphore   chan struct{}
	wg          sync.WaitGroup
}

// newDeploymentEnricher creates a new enricher with max concurrency limit.
func newDeploymentEnricher(
	imageClient v1.ImageServiceClient,
	cveName string,
	maxConcurrency int,
) *deploymentEnricher {
	return &deploymentEnricher{
		imageClient: imageClient,
		cveName:     cveName,
		semaphore:   make(chan struct{}, maxConcurrency),
	}
}

// enrich enriches a single deployment result with image data in a goroutine.
// Must be called before wait().
func (e *deploymentEnricher) enrich(
	ctx context.Context,
	deployment *DeploymentResult,
) {
	e.wg.Go(func() {
		e.semaphore <- struct{}{}

		defer func() { <-e.semaphore }()

		// Enrich the result in-place.
		images, err := fetchImagesForDeployment(ctx, e.imageClient, deployment, e.cveName)
		if err != nil {
			deployment.ImageFetchError = err.Error()

			return
		}

		deployment.AffectedImages = images
	})
}

// wait waits for all enrichment workers to complete.
func (e *deploymentEnricher) wait() {
	e.wg.Wait()
}

// fetchImagesForDeployment fetches images for a single deployment.
// It queries the images API filtered by CVE and Deployment ID.
func fetchImagesForDeployment(
	ctx context.Context,
	imageClient v1.ImageServiceClient,
	deployment *DeploymentResult,
	cveName string,
) ([]string, error) {
	query := fmt.Sprintf("CVE:%q+Deployment ID:%q", cveName, deployment.id)

	resp, err := imageClient.ListImages(ctx, &v1.RawQuery{Query: query})
	if err != nil {
		return nil, errors.Wrapf(err, "failed to fetch images for deployment %q in namespace %q",
			deployment.Name, deployment.Namespace)
	}

	images := make([]string, 0, len(resp.GetImages()))
	for _, img := range resp.GetImages() {
		images = append(images, img.GetName())
	}

	return images, nil
}

// handle is the handler for get_deployments_for_cve tool.
//
//nolint:funlen
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
	}

	currCursor, err := getCursor(&input)
	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: currCursor.GetOffset(),
			Limit:  defaultLimit + 1,
		},
	}

	resp, err := deploymentClient.ListDeployments(callCtx, listReq)
	if err != nil {
		return nil, nil, client.NewError(err, "ListDeployments")
	}

	rawDeployments := resp.GetDeployments()

	deployments := make([]DeploymentResult, len(rawDeployments))
	for i, deployment := range rawDeployments {
		deployments[i] = DeploymentResult{
			id:          deployment.GetId(),
			Name:        deployment.GetName(),
			Namespace:   deployment.GetNamespace(),
			ClusterID:   deployment.GetClusterId(),
			ClusterName: deployment.GetCluster(),
		}
	}

	if input.IncludeAffectedImages {
		imageClient := v1.NewImageServiceClient(conn)
		enricher := newDeploymentEnricher(imageClient, input.CVEName, defaultMaxFetchImageConcurrency)

		for i := range deployments {
			enricher.enrich(callCtx, &deployments[i])
		}

		enricher.wait()
	}

	// We always fetch limit+1 - if we do not have one additional element we can end paging.
	if len(deployments) <= defaultLimit {
		return nil, &getDeploymentsForCVEOutput{Deployments: deployments}, nil
	}

	nextCursorStr, err := currCursor.GetNextCursor(defaultLimit).Encode()
	if err != nil {
		return nil, nil, errors.Wrap(err, "unable to create next cursor")
	}

	output := &getDeploymentsForCVEOutput{
		Deployments: deployments[:len(deployments)-1],
		NextCursor:  nextCursorStr,
	}

	return nil, output, nil
}