extract pod selector labels

armosec · Jan 3, 2024 · cb9c404 · cb9c404
1 parent fb8ffc6
commit cb9c404
Show file tree

Hide file tree

Showing 2 changed files with 85 additions and 40 deletions.
diff --git a/armometadata/k8sutils.go b/armometadata/k8sutils.go
@@ -106,7 +106,7 @@ func LoadConfig(configPath string) (*ClusterConfig, error) {
 }
 
 // ExtractMetadataFromBytes extracts metadata from the JSON bytes of a Kubernetes object
-func ExtractMetadataFromJsonBytes(input []byte) (error, map[string]string, map[string]string, map[string]string, string, string, string, string) {
+func ExtractMetadataFromJsonBytes(input []byte) (error, map[string]string, map[string]string, map[string]string, string, string, string, string, map[string]string) {
 	// output values
 	annotations := map[string]string{}
 	labels := map[string]string{}
@@ -115,8 +115,9 @@ func ExtractMetadataFromJsonBytes(input []byte) (error, map[string]string, map[s
 	resourceVersion := ""
 	kind := ""
 	apiVersion := ""
+	podSelectorMatchLabels := map[string]string{}
 	// ujson parsing
-	var parent string
+	var parent, subParent, subParent2 string
 	err := ujson.Walk(input, func(level int, key, value []byte) bool {
 		switch level {
 		case 1:
@@ -129,39 +130,54 @@ func ExtractMetadataFromJsonBytes(input []byte) (error, map[string]string, map[s
 			}
 
 			// skip everything except metadata
-			if !bytes.EqualFold(key, []byte(`"metadata"`)) {
+			if !bytes.EqualFold(key, []byte(`"metadata"`)) && !bytes.EqualFold(key, []byte(`"spec"`)) {
 				return false
 			}
+
+			parent = unquote(key)
 		case 2:
-			// read creationTimestamp
-			if bytes.EqualFold(key, []byte(`"creationTimestamp"`)) {
-				creationTs = unquote(value)
-			}
-			// read resourceVersion
-			if bytes.EqualFold(key, []byte(`"resourceVersion"`)) {
-				resourceVersion = unquote(value)
+			if parent == "metadata" {
+				// read creationTimestamp
+				if bytes.EqualFold(key, []byte(`"creationTimestamp"`)) {
+					creationTs = unquote(value)
+				}
+				// read resourceVersion
+				if bytes.EqualFold(key, []byte(`"resourceVersion"`)) {
+					resourceVersion = unquote(value)
+				}
+
 			}
+
 			// record parent for level 3
-			parent = unquote(key)
+			subParent = unquote(key)
+
 		case 3:
 			// read annotations
-			if parent == "annotations" {
+			if subParent == "annotations" {
 				annotations[unquote(key)] = unquote(value)
 			}
 			// read labels
-			if parent == "labels" {
+			if subParent == "labels" {
 				labels[unquote(key)] = unquote(value)
 			}
+
+			subParent2 = unquote(key)
+
 		case 4:
 			// read ownerReferences
-			if parent == "ownerReferences" {
+			if subParent == "ownerReferences" {
 				ownerReferences[unquote(key)] = unquote(value)
 			}
 
+			if subParent2 == "matchLabels" {
+				podSelectorMatchLabels[unquote(key)] = unquote(value)
+
+			}
+
 		}
 		return true
 	})
-	return err, annotations, labels, ownerReferences, creationTs, resourceVersion, kind, apiVersion
+	return err, annotations, labels, ownerReferences, creationTs, resourceVersion, kind, apiVersion, podSelectorMatchLabels
 }
 
 func unquote(value []byte) string {

diff --git a/armometadata/k8sutils_test.go b/armometadata/k8sutils_test.go
@@ -124,16 +124,41 @@ func BoolPtr(b bool) *bool {
 
 func TestExtractMetadataFromJsonBytes(t *testing.T) {
 	tests := []struct {
-		name            string
-		want            error
-		annotations     map[string]string
-		labels          map[string]string
-		ownerReferences map[string]string
-		creationTs      string
-		resourceVersion string
-		kind            string
-		apiVersion      string
+		name                   string
+		want                   error
+		annotations            map[string]string
+		labels                 map[string]string
+		ownerReferences        map[string]string
+		creationTs             string
+		resourceVersion        string
+		kind                   string
+		apiVersion             string
+		podSelectorMatchLabels map[string]string
 	}{
+		{
+			name:                   "networkpolicy_withoutmatching_labels",
+			annotations:            map[string]string{},
+			labels:                 map[string]string{},
+			ownerReferences:        map[string]string{},
+			creationTs:             "2023-11-16T10:12:35Z",
+			resourceVersion:        "",
+			kind:                   "NetworkPolicy",
+			apiVersion:             "networking.k8s.io/v1",
+			podSelectorMatchLabels: map[string]string{},
+		},
+		{
+			name:            "networkpolicy_withmatching_labels",
+			annotations:     map[string]string{},
+			labels:          map[string]string{},
+			ownerReferences: map[string]string{},
+			creationTs:      "2023-11-16T10:12:35Z",
+			resourceVersion: "",
+			kind:            "NetworkPolicy",
+			apiVersion:      "networking.k8s.io/v1",
+			podSelectorMatchLabels: map[string]string{
+				"role": "frontend",
+			},
+		},
 		{
 			name: "applicationactivity",
 			annotations: map[string]string{
@@ -147,11 +172,12 @@ func TestExtractMetadataFromJsonBytes(t *testing.T) {
 				"kubescape.io/workload-name":        "storage",
 				"kubescape.io/workload-namespace":   "kubescape",
 			},
-			ownerReferences: map[string]string{},
-			creationTs:      "2023-11-16T10:15:05Z",
-			resourceVersion: "1",
-			kind:            "ApplicationActivity",
-			apiVersion:      "spdx.softwarecomposition.kubescape.io/v1beta1",
+			ownerReferences:        map[string]string{},
+			creationTs:             "2023-11-16T10:15:05Z",
+			resourceVersion:        "1",
+			kind:                   "ApplicationActivity",
+			apiVersion:             "spdx.softwarecomposition.kubescape.io/v1beta1",
+			podSelectorMatchLabels: map[string]string{},
 		},
 		{
 			name: "pod",
@@ -178,10 +204,11 @@ func TestExtractMetadataFromJsonBytes(t *testing.T) {
 				"name":               "kubescape-549f95c69",
 				"uid":                "c0ff7d3b-4183-482c-81c5-998faf0b6150",
 			},
-			creationTs:      "2023-11-16T10:12:35Z",
-			resourceVersion: "59348379",
-			kind:            "Pod",
-			apiVersion:      "v1",
+			creationTs:             "2023-11-16T10:12:35Z",
+			resourceVersion:        "59348379",
+			kind:                   "Pod",
+			apiVersion:             "v1",
+			podSelectorMatchLabels: map[string]string{},
 		},
 		{
 			name: "sbom",
@@ -193,18 +220,19 @@ func TestExtractMetadataFromJsonBytes(t *testing.T) {
 				"kubescape.io/image-id":   "quay-io-kubescape-kubescape-sha256-608b85d3de51caad84a2bfe089ec",
 				"kubescape.io/image-name": "quay-io-kubescape-kubescape",
 			},
-			ownerReferences: map[string]string{},
-			creationTs:      "2023-11-16T10:13:40Z",
-			resourceVersion: "1",
-			kind:            "SBOMSPDXv2p3",
-			apiVersion:      "spdx.softwarecomposition.kubescape.io/v1beta1",
+			ownerReferences:        map[string]string{},
+			creationTs:             "2023-11-16T10:13:40Z",
+			resourceVersion:        "1",
+			kind:                   "SBOMSPDXv2p3",
+			apiVersion:             "spdx.softwarecomposition.kubescape.io/v1beta1",
+			podSelectorMatchLabels: map[string]string{},
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			input, err := os.ReadFile(fmt.Sprintf("testdata/%s.json", tt.name))
 			assert.NoError(t, err)
-			got, annotations, labels, ownerReferences, creationTs, resourceVersion, kind, apiVersion := ExtractMetadataFromJsonBytes(input)
+			got, annotations, labels, ownerReferences, creationTs, resourceVersion, kind, apiVersion, podSelectorMatchLabels := ExtractMetadataFromJsonBytes(input)
 			assert.Equal(t, tt.want, got)
 			assert.Equal(t, tt.annotations, annotations)
 			assert.Equal(t, tt.labels, labels)
@@ -213,6 +241,7 @@ func TestExtractMetadataFromJsonBytes(t *testing.T) {
 			assert.Equal(t, tt.resourceVersion, resourceVersion)
 			assert.Equal(t, tt.kind, kind)
 			assert.Equal(t, tt.apiVersion, apiVersion)
+			assert.Equal(t, tt.podSelectorMatchLabels, podSelectorMatchLabels)
 		})
 	}
 }
@@ -221,6 +250,6 @@ func BenchmarkExtractMetadataFromJsonBytes(b *testing.B) {
 	input, err := os.ReadFile("testdata/applicationactivity.json")
 	assert.NoError(b, err)
 	for i := 0; i < b.N; i++ {
-		_, _, _, _, _, _, _, _ = ExtractMetadataFromJsonBytes(input)
+		_, _, _, _, _, _, _, _, _ = ExtractMetadataFromJsonBytes(input)
 	}
 }