add struct for metadata

Signed-off-by: Matthias Bertschy <[email protected]>
armosec · Jan 3, 2024 · 5b9c1ab · 5b9c1ab
1 parent 3a908c5
commit 5b9c1ab
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 32 deletions.
diff --git a/armometadata/k8sutils.go b/armometadata/k8sutils.go
@@ -105,31 +105,40 @@ func LoadConfig(configPath string) (*ClusterConfig, error) {
 	return config, err
 }
 
+type Metadata struct {
+	Annotations            map[string]string
+	Labels                 map[string]string
+	OwnerReferences        map[string]string
+	CreationTimestamp      string
+	ResourceVersion        string
+	Kind                   string
+	ApiVersion             string
+	PodSelectorMatchLabels map[string]string
+}
+
 // 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, map[string]string) {
+func ExtractMetadataFromJsonBytes(input []byte) (Metadata, error) {
 	// output values
-	annotations := map[string]string{}
-	labels := map[string]string{}
-	ownerReferences := map[string]string{}
-	creationTs := ""
-	resourceVersion := ""
-	kind := ""
-	apiVersion := ""
-	podSelectorMatchLabels := map[string]string{}
+	m := Metadata{
+		Annotations:            map[string]string{},
+		Labels:                 map[string]string{},
+		OwnerReferences:        map[string]string{},
+		PodSelectorMatchLabels: map[string]string{},
+	}
 	// ujson parsing
 	var parent, subParent, subParent2 string
 	err := ujson.Walk(input, func(level int, key, value []byte) bool {
 		switch level {
 		case 1:
 			if bytes.EqualFold(key, []byte(`"kind"`)) {
-				kind = unquote(value)
+				m.Kind = unquote(value)
 			}
 
 			if bytes.EqualFold(key, []byte(`"apiVersion"`)) {
-				apiVersion = unquote(value)
+				m.ApiVersion = unquote(value)
 			}
 
-			// skip everything except metadata
+			// skip everything except metadata and spec
 			if !bytes.EqualFold(key, []byte(`"metadata"`)) && !bytes.EqualFold(key, []byte(`"spec"`)) {
 				return false
 			}
@@ -139,11 +148,11 @@ func ExtractMetadataFromJsonBytes(input []byte) (error, map[string]string, map[s
 			if parent == "metadata" {
 				// read creationTimestamp
 				if bytes.EqualFold(key, []byte(`"creationTimestamp"`)) {
-					creationTs = unquote(value)
+					m.CreationTimestamp = unquote(value)
 				}
 				// read resourceVersion
 				if bytes.EqualFold(key, []byte(`"resourceVersion"`)) {
-					resourceVersion = unquote(value)
+					m.ResourceVersion = unquote(value)
 				}
 
 			}
@@ -154,30 +163,29 @@ func ExtractMetadataFromJsonBytes(input []byte) (error, map[string]string, map[s
 		case 3:
 			// read annotations
 			if subParent == "annotations" {
-				annotations[unquote(key)] = unquote(value)
+				m.Annotations[unquote(key)] = unquote(value)
 			}
 			// read labels
 			if subParent == "labels" {
-				labels[unquote(key)] = unquote(value)
+				m.Labels[unquote(key)] = unquote(value)
 			}
 
 			subParent2 = unquote(key)
 
 		case 4:
 			// read ownerReferences
 			if subParent == "ownerReferences" {
-				ownerReferences[unquote(key)] = unquote(value)
+				m.OwnerReferences[unquote(key)] = unquote(value)
 			}
 
 			if subParent2 == "matchLabels" {
-				podSelectorMatchLabels[unquote(key)] = unquote(value)
-
+				m.PodSelectorMatchLabels[unquote(key)] = unquote(value)
 			}
 
 		}
 		return true
 	})
-	return err, annotations, labels, ownerReferences, creationTs, resourceVersion, kind, apiVersion, podSelectorMatchLabels
+	return m, err
 }
 
 func unquote(value []byte) string {

diff --git a/armometadata/k8sutils_test.go b/armometadata/k8sutils_test.go
@@ -125,7 +125,7 @@ func BoolPtr(b bool) *bool {
 func TestExtractMetadataFromJsonBytes(t *testing.T) {
 	tests := []struct {
 		name                   string
-		want                   error
+		wantErr                error
 		annotations            map[string]string
 		labels                 map[string]string
 		ownerReferences        map[string]string
@@ -232,16 +232,16 @@ func TestExtractMetadataFromJsonBytes(t *testing.T) {
 		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, podSelectorMatchLabels := ExtractMetadataFromJsonBytes(input)
-			assert.Equal(t, tt.want, got)
-			assert.Equal(t, tt.annotations, annotations)
-			assert.Equal(t, tt.labels, labels)
-			assert.Equal(t, tt.ownerReferences, ownerReferences)
-			assert.Equal(t, tt.creationTs, creationTs)
-			assert.Equal(t, tt.resourceVersion, resourceVersion)
-			assert.Equal(t, tt.kind, kind)
-			assert.Equal(t, tt.apiVersion, apiVersion)
-			assert.Equal(t, tt.podSelectorMatchLabels, podSelectorMatchLabels)
+			m, err := ExtractMetadataFromJsonBytes(input)
+			assert.Equal(t, tt.wantErr, err)
+			assert.Equal(t, tt.annotations, m.Annotations)
+			assert.Equal(t, tt.labels, m.Labels)
+			assert.Equal(t, tt.ownerReferences, m.OwnerReferences)
+			assert.Equal(t, tt.creationTs, m.CreationTimestamp)
+			assert.Equal(t, tt.resourceVersion, m.ResourceVersion)
+			assert.Equal(t, tt.kind, m.Kind)
+			assert.Equal(t, tt.apiVersion, m.ApiVersion)
+			assert.Equal(t, tt.podSelectorMatchLabels, m.PodSelectorMatchLabels)
 		})
 	}
 }
@@ -250,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)
 	}
 }