From 6594dd82bba6163e34fe2c63a5faff86dafe804f Mon Sep 17 00:00:00 2001
From: kooomix <eranm@armosec.io>
Date: Mon, 27 Nov 2023 08:44:04 +0200
Subject: [PATCH] extend ExtractMetadataFromJsonBytes

---
 armometadata/k8sutils.go      | 14 ++++++++++++--
 armometadata/k8sutils_test.go | 14 ++++++++++++--
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/armometadata/k8sutils.go b/armometadata/k8sutils.go
index c9481f2..15b9476 100644
--- a/armometadata/k8sutils.go
+++ b/armometadata/k8sutils.go
@@ -106,18 +106,28 @@ 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) {
+func ExtractMetadataFromJsonBytes(input []byte) (error, map[string]string, map[string]string, map[string]string, string, string, string, string) {
 	// output values
 	annotations := map[string]string{}
 	labels := map[string]string{}
 	ownerReferences := map[string]string{}
 	creationTs := ""
 	resourceVersion := ""
+	kind := ""
+	apiVersion := ""
 	// ujson parsing
 	var parent 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)
+			}
+
+			if bytes.EqualFold(key, []byte(`"apiVersion"`)) {
+				apiVersion = unquote(value)
+			}
+
 			// skip everything except metadata
 			if !bytes.EqualFold(key, []byte(`"metadata"`)) {
 				return false
@@ -151,7 +161,7 @@ func ExtractMetadataFromJsonBytes(input []byte) (error, map[string]string, map[s
 		}
 		return true
 	})
-	return err, annotations, labels, ownerReferences, creationTs, resourceVersion
+	return err, annotations, labels, ownerReferences, creationTs, resourceVersion, kind, apiVersion
 }
 
 func unquote(value []byte) string {
diff --git a/armometadata/k8sutils_test.go b/armometadata/k8sutils_test.go
index ec3a96b..0e4cab3 100644
--- a/armometadata/k8sutils_test.go
+++ b/armometadata/k8sutils_test.go
@@ -131,6 +131,8 @@ func TestExtractMetadataFromJsonBytes(t *testing.T) {
 		ownerReferences map[string]string
 		creationTs      string
 		resourceVersion string
+		kind            string
+		apiVersion      string
 	}{
 		{
 			name: "applicationactivity",
@@ -148,6 +150,8 @@ func TestExtractMetadataFromJsonBytes(t *testing.T) {
 			ownerReferences: map[string]string{},
 			creationTs:      "2023-11-16T10:15:05Z",
 			resourceVersion: "1",
+			kind:            "ApplicationActivity",
+			apiVersion:      "spdx.softwarecomposition.kubescape.io/v1beta1",
 		},
 		{
 			name: "pod",
@@ -176,6 +180,8 @@ func TestExtractMetadataFromJsonBytes(t *testing.T) {
 			},
 			creationTs:      "2023-11-16T10:12:35Z",
 			resourceVersion: "59348379",
+			kind:            "Pod",
+			apiVersion:      "v1",
 		},
 		{
 			name: "sbom",
@@ -190,19 +196,23 @@ func TestExtractMetadataFromJsonBytes(t *testing.T) {
 			ownerReferences: map[string]string{},
 			creationTs:      "2023-11-16T10:13:40Z",
 			resourceVersion: "1",
+			kind:            "SBOMSPDXv2p3",
+			apiVersion:      "spdx.softwarecomposition.kubescape.io/v1beta1",
 		},
 	}
 	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 := ExtractMetadataFromJsonBytes(input)
+			got, annotations, labels, ownerReferences, creationTs, resourceVersion, kind, apiVersion := 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)
 		})
 	}
 }
@@ -211,6 +221,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)
 	}
 }