Merge pull request #1915 from majiayu000/fix-1660-define-labels-to-copy-from-wor-1231-0603

feat: Define labels to copy from workloads to reports

cmd/scan/scan.go

@@ -93,6 +93,7 @@ func GetScanCommand(ks meta.IKubescape) *cobra.Command {
 	scanCmd.PersistentFlags().BoolVarP(&scanInfo.EnableRegoPrint, "enable-rego-prints", "", false, "Enable sending to rego prints to the logs (use with debug log level: -l debug)")
 	scanCmd.PersistentFlags().BoolVarP(&scanInfo.ScanImages, "scan-images", "", false, "Scan resources images")
 	scanCmd.PersistentFlags().BoolVarP(&scanInfo.UseDefaultMatchers, "use-default-matchers", "", true, "Use default matchers (true) or CPE matchers (false) for image scanning")
+	scanCmd.PersistentFlags().StringSliceVar(&scanInfo.LabelsToCopy, "labels-to-copy", nil, "Labels to copy from workloads to scan reports for easy identification. e.g: --labels-to-copy=app,team,environment")
 
 	scanCmd.PersistentFlags().MarkDeprecated("fail-threshold", "use '--compliance-threshold' flag instead. Flag will be removed at 1.Dec.2023")
 	scanCmd.PersistentFlags().MarkDeprecated("create-account", "Create account is no longer supported. In case of a missing Account ID and a configured backend server, a new account id will be generated automatically by Kubescape. Feel free to contact the Kubescape maintainers for more information.")

core/cautils/datastructures.go

@@ -69,6 +69,7 @@ type OPASessionObj struct {
 	TopWorkloadsByScore   []reporthandling.IResource
 	TemplateMapping       map[string]MappingNodes // Map chart obj to template (only for rendering from path)
 	TriggeredByCLI        bool
+	LabelsToCopy          []string // Labels to copy from workloads to scan reports
 }
 
 func NewOPASessionObj(ctx context.Context, frameworks []reporthandling.Framework, k8sResources K8SResources, scanInfo *ScanInfo) *OPASessionObj {
@@ -87,6 +88,7 @@ func NewOPASessionObj(ctx context.Context, frameworks []reporthandling.Framework
 		OmitRawResources:      scanInfo.OmitRawResources,
 		TriggeredByCLI:        scanInfo.TriggeredByCLI,
 		TemplateMapping:       make(map[string]MappingNodes),
+		LabelsToCopy:          scanInfo.LabelsToCopy,
 	}
 }
 

core/cautils/scaninfo.go

@@ -140,6 +140,7 @@ type ScanInfo struct {
 	UseDefaultMatchers    bool
 	ChartPath             string
 	FilePath              string
+	LabelsToCopy          []string // Labels to copy from workloads to scan reports
 	scanningContext       *ScanningContext
 	cleanups              []func()
 }

core/pkg/resultshandling/printer/v2/jsonprinter.go

@@ -121,8 +121,9 @@ func printConfigurationsScanning(opaSessionObj *cautils.OPASessionObj, imageScan
 	}
 
 	// Convert to PostureReportWithSeverity to add severity field to controls
+	// and extract specified labels from workloads
 	finalizedReport := FinalizeResults(opaSessionObj)
-	reportWithSeverity := ConvertToPostureReportWithSeverity(finalizedReport)
+	reportWithSeverity := ConvertToPostureReportWithSeverityAndLabels(finalizedReport, opaSessionObj.LabelsToCopy, opaSessionObj.AllResources)
 
 	r, err := json.Marshal(reportWithSeverity)
 	_, err = jp.writer.Write(r)

core/pkg/resultshandling/printer/v2/utils.go

@@ -61,6 +61,7 @@ type PostureReportWithSeverity struct {
 	Attributes           []reportsummary.PostureAttributes `json:"attributes"`
 	Results              []ResultWithSeverity              `json:"results,omitempty"`
 	Metadata             reporthandlingv2.Metadata         `json:"metadata,omitempty"`
+	ResourceLabels       map[string]map[string]string      `json:"resourceLabels,omitempty"` // map[resourceID]map[labelKey]labelValue - extracted labels from workloads
 }
 
 // enrichControlsWithSeverity adds severity field to controls based on scoreFactor
@@ -103,12 +104,24 @@ func enrichResultsWithSeverity(results []resourcesresults.Result, controlSummari
 
 // ConvertToPostureReportWithSeverity converts PostureReport to PostureReportWithSeverity
 func ConvertToPostureReportWithSeverity(report *reporthandlingv2.PostureReport) *PostureReportWithSeverity {
+	return ConvertToPostureReportWithSeverityAndLabels(report, nil, nil)
+}
+
+// ConvertToPostureReportWithSeverityAndLabels converts PostureReport to PostureReportWithSeverity
+// and extracts specified labels from workloads
+func ConvertToPostureReportWithSeverityAndLabels(report *reporthandlingv2.PostureReport, labelsToCopy []string, allResources map[string]workloadinterface.IMetadata) *PostureReportWithSeverity {
 	if report == nil {
 		return nil
 	}
 	enrichedControls := enrichControlsWithSeverity(report.SummaryDetails.Controls)
 	enrichedResults := enrichResultsWithSeverity(report.Results, report.SummaryDetails.Controls)
 
+	// Extract labels from resources if labelsToCopy is specified
+	var resourceLabels map[string]map[string]string
+	if len(labelsToCopy) > 0 && allResources != nil {
+		resourceLabels = extractResourceLabels(allResources, labelsToCopy)
+	}
+
 	return &PostureReportWithSeverity{
 		ReportGenerationTime: report.ReportGenerationTime.Format("2006-01-02T15:04:05Z07:00"),
 		ClusterAPIServerInfo: report.ClusterAPIServerInfo,
@@ -126,13 +139,46 @@ func ConvertToPostureReportWithSeverity(report *reporthandlingv2.PostureReport)
 			Score:                     report.SummaryDetails.Score,
 			ComplianceScore:           report.SummaryDetails.ComplianceScore,
 		},
-		Resources:  report.Resources,
-		Attributes: report.Attributes,
-		Results:    enrichedResults,
-		Metadata:   report.Metadata,
+		Resources:      report.Resources,
+		Attributes:     report.Attributes,
+		Results:        enrichedResults,
+		Metadata:       report.Metadata,
+		ResourceLabels: resourceLabels,
 	}
 }
 
+// extractResourceLabels extracts specified labels from all resources
+func extractResourceLabels(allResources map[string]workloadinterface.IMetadata, labelsToCopy []string) map[string]map[string]string {
+	resourceLabels := make(map[string]map[string]string)
+
+	for resourceID, resource := range allResources {
+		// IMetadata doesn't have GetLabels, need to cast to IBasicWorkload
+		basicWorkload, ok := resource.(workloadinterface.IBasicWorkload)
+		if !ok {
+			continue
+		}
+
+		labels := basicWorkload.GetLabels()
+		if labels == nil {
+			continue
+		}
+
+		extractedLabels := make(map[string]string)
+		for _, labelKey := range labelsToCopy {
+			if value, exists := labels[labelKey]; exists {
+				extractedLabels[labelKey] = value
+			}
+		}
+
+		// Only add to result if at least one label was found
+		if len(extractedLabels) > 0 {
+			resourceLabels[resourceID] = extractedLabels
+		}
+	}
+
+	return resourceLabels
+}
+
 // FinalizeResults finalize the results objects by copying data from map to lists
 func FinalizeResults(data *cautils.OPASessionObj) *reporthandlingv2.PostureReport {
 	report := reporthandlingv2.PostureReport{

core/pkg/resultshandling/printer/v2/utils_test.go

@@ -1,12 +1,14 @@
 package printer
 
 import (
+	"encoding/json"
 	"testing"
 
 	v5 "github.com/anchore/grype/grype/db/v5"
 	"github.com/anchore/grype/grype/match"
 	"github.com/anchore/grype/grype/pkg"
 	"github.com/anchore/grype/grype/vulnerability"
+	"github.com/kubescape/k8s-interface/workloadinterface"
 	"github.com/kubescape/kubescape/v3/core/pkg/resultshandling/printer/v2/prettyprinter/tableprinter/imageprinter"
 	"github.com/stretchr/testify/assert"
 )
@@ -742,3 +744,131 @@ func TestSetSeverityToSummaryMap(t *testing.T) {
 		})
 	}
 }
+
+func createWorkloadWithLabels(name, namespace string, labels map[string]string) workloadinterface.IMetadata {
+	// Convert labels to map[string]interface{} for JSON marshaling
+	labelsInterface := make(map[string]interface{})
+	for k, v := range labels {
+		labelsInterface[k] = v
+	}
+
+	obj := map[string]interface{}{
+		"apiVersion": "apps/v1",
+		"kind":       "Deployment",
+		"metadata": map[string]interface{}{
+			"name":      name,
+			"namespace": namespace,
+			"labels":    labelsInterface,
+		},
+		"spec": map[string]interface{}{
+			"template": map[string]interface{}{
+				"spec": map[string]interface{}{
+					"containers": []interface{}{},
+				},
+			},
+		},
+	}
+	objBytes, _ := json.Marshal(obj)
+	workload, _ := workloadinterface.NewWorkload(objBytes)
+	return workload
+}
+
+func TestExtractResourceLabels(t *testing.T) {
+	tests := []struct {
+		name         string
+		allResources map[string]workloadinterface.IMetadata
+		labelsToCopy []string
+		want         map[string]map[string]string
+	}{
+		{
+			name:         "empty resources",
+			allResources: map[string]workloadinterface.IMetadata{},
+			labelsToCopy: []string{"app", "team"},
+			want:         map[string]map[string]string{},
+		},
+		{
+			name:         "empty labels to copy",
+			allResources: map[string]workloadinterface.IMetadata{},
+			labelsToCopy: []string{},
+			want:         map[string]map[string]string{},
+		},
+		{
+			name: "single resource with matching labels",
+			allResources: map[string]workloadinterface.IMetadata{
+				"resource-1": createWorkloadWithLabels("test-deploy", "default", map[string]string{
+					"app":     "myapp",
+					"team":    "platform",
+					"version": "v1",
+				}),
+			},
+			labelsToCopy: []string{"app", "team"},
+			want: map[string]map[string]string{
+				"resource-1": {
+					"app":  "myapp",
+					"team": "platform",
+				},
+			},
+		},
+		{
+			name: "single resource with partial matching labels",
+			allResources: map[string]workloadinterface.IMetadata{
+				"resource-1": createWorkloadWithLabels("test-deploy", "default", map[string]string{
+					"app": "myapp",
+				}),
+			},
+			labelsToCopy: []string{"app", "team"},
+			want: map[string]map[string]string{
+				"resource-1": {
+					"app": "myapp",
+				},
+			},
+		},
+		{
+			name: "single resource with no matching labels",
+			allResources: map[string]workloadinterface.IMetadata{
+				"resource-1": createWorkloadWithLabels("test-deploy", "default", map[string]string{
+					"version": "v1",
+				}),
+			},
+			labelsToCopy: []string{"app", "team"},
+			want:         map[string]map[string]string{},
+		},
+		{
+			name: "multiple resources with various labels",
+			allResources: map[string]workloadinterface.IMetadata{
+				"resource-1": createWorkloadWithLabels("deploy-1", "default", map[string]string{
+					"app":  "app1",
+					"team": "team1",
+				}),
+				"resource-2": createWorkloadWithLabels("deploy-2", "default", map[string]string{
+					"app": "app2",
+				}),
+				"resource-3": createWorkloadWithLabels("deploy-3", "default", map[string]string{
+					"version": "v1",
+				}),
+			},
+			labelsToCopy: []string{"app", "team"},
+			want: map[string]map[string]string{
+				"resource-1": {
+					"app":  "app1",
+					"team": "team1",
+				},
+				"resource-2": {
+					"app": "app2",
+				},
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := extractResourceLabels(tt.allResources, tt.labelsToCopy)
+			assert.Equal(t, len(tt.want), len(got), "number of resources with extracted labels should match")
+			for resourceID, wantLabels := range tt.want {
+				gotLabels, ok := got[resourceID]
+				assert.True(t, ok, "resource %s should be present in result", resourceID)
+				assert.Equal(t, wantLabels, gotLabels, "labels for resource %s should match", resourceID)
+			}
+		})
+	}
+}