From c3604c1170f346c93060f3af0de2fd087f675760 Mon Sep 17 00:00:00 2001
From: Sambhav Kothari <sambhavs.email@gmail.com>
Date: Sat, 7 May 2022 13:05:04 +0100
Subject: [PATCH] Add an object_from_lists function (#3824)

---
 pkg/engine/jmespath/functions.go           | 39 ++++++++++++++++++
 pkg/engine/jmespath/functions_test.go      | 47 ++++++++++++++++++++++
 test/cli/test-mutate/kyverno-test.yaml     |  6 +++
 test/cli/test-mutate/patched-resource.yaml |  7 ++++
 test/cli/test-mutate/policy.yaml           | 24 +++++++++++
 test/cli/test-mutate/resource.yaml         |  9 +++++
 6 files changed, 132 insertions(+)
 create mode 100644 test/cli/test-mutate/patched-resource.yaml

diff --git a/pkg/engine/jmespath/functions.go b/pkg/engine/jmespath/functions.go
index 16dd46f121..8651b97fb1 100644
--- a/pkg/engine/jmespath/functions.go
+++ b/pkg/engine/jmespath/functions.go
@@ -64,6 +64,7 @@ var (
 	parseJson              = "parse_json"
 	parseYAML              = "parse_yaml"
 	items                  = "items"
+	objectFromLists        = "object_from_lists"
 )
 
 const errorPrefix = "JMESPath function '%s': "
@@ -382,6 +383,17 @@ func GetFunctions() []*FunctionEntry {
 			ReturnType: []JpType{JpArray},
 			Note:       "converts a map to an array of objects where each key:value is an item in the array",
 		},
+		{
+			Entry: &gojmespath.FunctionEntry{Name: objectFromLists,
+				Arguments: []ArgSpec{
+					{Types: []JpType{JpArray}},
+					{Types: []JpType{JpArray}},
+				},
+				Handler: jpObjectFromLists,
+			},
+			ReturnType: []JpType{JpObject},
+			Note:       "converts a pair of lists containing keys and values to an object",
+		},
 	}
 
 }
@@ -844,6 +856,33 @@ func jpItems(arguments []interface{}) (interface{}, error) {
 	return arrayOfObj, nil
 }
 
+func jpObjectFromLists(arguments []interface{}) (interface{}, error) {
+	keys, ok := arguments[0].([]interface{})
+	if !ok {
+		return nil, fmt.Errorf(invalidArgumentTypeError, arguments, 0, "Array")
+	}
+	values, ok := arguments[1].([]interface{})
+	if !ok {
+		return nil, fmt.Errorf(invalidArgumentTypeError, arguments, 1, "Array")
+	}
+
+	output := map[string]interface{}{}
+
+	for i, ikey := range keys {
+		key, err := ifaceToString(ikey)
+		if err != nil {
+			return nil, fmt.Errorf(invalidArgumentTypeError, arguments, 0, "StringArray")
+		}
+		if i < len(values) {
+			output[key] = values[i]
+		} else {
+			output[key] = nil
+		}
+	}
+
+	return output, nil
+}
+
 // InterfaceToString casts an interface to a string type
 func ifaceToString(iface interface{}) (string, error) {
 	switch i := iface.(type) {
diff --git a/pkg/engine/jmespath/functions_test.go b/pkg/engine/jmespath/functions_test.go
index 5b421827cf..7b63507e52 100644
--- a/pkg/engine/jmespath/functions_test.go
+++ b/pkg/engine/jmespath/functions_test.go
@@ -1334,3 +1334,50 @@ func Test_Items(t *testing.T) {
 	}
 
 }
+
+func Test_ObjectFromLists(t *testing.T) {
+
+	testCases := []struct {
+		keys           string
+		values         string
+		expectedResult map[string]interface{}
+	}{
+		{
+			keys:   `["key1", "key2"]`,
+			values: `["1", "2"]`,
+			expectedResult: map[string]interface{}{
+				"key1": "1",
+				"key2": "2",
+			},
+		},
+		{
+			keys:   `["key1", "key2"]`,
+			values: `[1, "2"]`,
+			expectedResult: map[string]interface{}{
+				"key1": 1.0,
+				"key2": "2",
+			},
+		},
+		{
+			keys:   `["key1", "key2"]`,
+			values: `[1]`,
+			expectedResult: map[string]interface{}{
+				"key1": 1.0,
+				"key2": nil,
+			},
+		},
+	}
+
+	for i, tc := range testCases {
+		t.Run(fmt.Sprintf("case %d", i), func(t *testing.T) {
+			query, err := New("object_from_lists(`" + tc.keys + "`,`" + tc.values + "`)")
+			assert.NilError(t, err)
+			res, err := query.Search("")
+			assert.NilError(t, err)
+			result, ok := res.(map[string]interface{})
+			assert.Assert(t, ok)
+			assert.DeepEqual(t, result, tc.expectedResult)
+		})
+	}
+
+}
diff --git a/test/cli/test-mutate/kyverno-test.yaml b/test/cli/test-mutate/kyverno-test.yaml
index e07ae443b7..ba829a7553 100644
--- a/test/cli/test-mutate/kyverno-test.yaml
+++ b/test/cli/test-mutate/kyverno-test.yaml
@@ -82,3 +82,9 @@ results:
     patchedResource: patchedResource11.yaml
     kind: Pod
     result: skip
+  - policy: example
+    rule: object_from_lists
+    resource: example
+    patchedResource: patched-resource.yaml
+    kind: Pod
+    result: pass
\ No newline at end of file
diff --git a/test/cli/test-mutate/patched-resource.yaml b/test/cli/test-mutate/patched-resource.yaml
new file mode 100644
index 0000000000..149bccc81a
--- /dev/null
+++ b/test/cli/test-mutate/patched-resource.yaml
@@ -0,0 +1,7 @@
+apiVersion: v1
+kind: Pod
+metadata:
+  name: example
+  annotations:
+    key1: "1"
+    key2: "1"
diff --git a/test/cli/test-mutate/policy.yaml b/test/cli/test-mutate/policy.yaml
index 767b5d6cd2..ce008db582 100644
--- a/test/cli/test-mutate/policy.yaml
+++ b/test/cli/test-mutate/policy.yaml
@@ -75,3 +75,27 @@ spec:
             options:
               - name: ndots
                 value: "1"
+---
+apiVersion: kyverno.io/v1
+kind: ClusterPolicy
+metadata:
+  name: example
+spec:
+  rules:
+  - name: object_from_lists
+    context:
+      - name: annotations
+        variable:
+          jmesPath: items(request.object.metadata.annotations, 'key', 'value')[?starts_with(key, 'key')]
+      - name: annotations
+        variable:
+          jmesPath: object_from_lists(annotations[].key, annotations[].value)
+    match:
+      resources:
+        kinds:
+        - Pod
+    mutate:
+      patchesJson6902: |-
+        - path: "/metadata/annotations"
+          op: replace
+          value: {{ annotations }}
diff --git a/test/cli/test-mutate/resource.yaml b/test/cli/test-mutate/resource.yaml
index b66186af9e..d49f5665d6 100644
--- a/test/cli/test-mutate/resource.yaml
+++ b/test/cli/test-mutate/resource.yaml
@@ -98,3 +98,12 @@ spec:
     image: nginx:latest
 
 
+---
+apiVersion: v1
+kind: Pod
+metadata:
+  name: example
+  annotations:
+    key1: "1"
+    key2: "1"
+    notkey: "2"