Merge pull request #50 from Esonhugh/main

helperfunctions to common Slice operations
projectdiscovery · Jun 28, 2023 · 71da38c · 71da38c
2 parents 6f7f739 + c8ee1b1
commit 71da38c
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 1 deletion.
diff --git a/README.md b/README.md
@@ -83,6 +83,7 @@ func main() {
 | hmac(algorithm, data, secret) string                                  | hmac function that accepts a hashing function type with data and secret                                             | `hmac("sha1", "test", "scrt")`                                                                                                                       | `8856b111056d946d5c6c92a21b43c233596623c6`                                                                                                                                                                                                                                                                                                                                                 |
 | html_escape(input interface{}) string                                 | HTML escapes the given input                                                                                        | `html_escape("<body>test</body>")`                                                                                                                   | `&lt;body&gt;test&lt;/body&gt;`                                                                                                                                                                                                                                                                                                                                                            |
 | html_unescape(input interface{}) string                               | HTML un-escapes the given input                                                                                     | `html_unescape("&lt;body&gt;test&lt;/body&gt;")`                                                                                                     | `<body>test</body>`                                                                                                                                                                                                                                                                                                                                                                        |
+| index(slice, index) interface{}                               | Select item at index from slice or string (zero based)                                                                                     | `index("test",0)`                                                                                                     | `t`                                                                                                                                                                                                                                                                                                                                                                        |
 | join(separator string, elements ...interface{}) string                | Joins the given elements using the specified separator                                                              | `join("_", 123, "hello", "world")`                                                                                                                   | `123_hello_world`                                                                                                                                                                                                                                                                                                                                                                          |
 | json_minify(json) string | Minifies a JSON string by removing unnecessary whitespace | `json_minify("{ \"name\": \"John Doe\", \"foo\": \"bar\" }")` | `{"foo":"bar","name":"John Doe"}` |
 | json_prettify(json) string | Prettifies a JSON string by adding indentation | `json_prettify("{\"foo\":\"bar\",\"name\":\"John Doe\"}")` | `{\n \"foo\": \"bar\",\n \"name\": \"John Doe\"\n}` |

diff --git a/dsl.go b/dsl.go
@@ -21,6 +21,7 @@ import (
 	"io"
 	"math"
 	"net/url"
+	"reflect"
 	"regexp"
 	"sort"
 	"strconv"
@@ -76,8 +77,42 @@ func MustAddFunction(function dslFunction) {
 }
 
 func init() {
+	// note: index helper is zero based
+	MustAddFunction(NewWithPositionalArgs("index", 2, func(args ...interface{}) (interface{}, error) {
+		index, err := strconv.ParseInt(toString(args[1]), 10, 64)
+		if err != nil {
+			return nil, err
+		}
+		// If the first argument is a slice, we index into it
+		switch v := args[0].(type) {
+		case []string:
+			l := int64(len(v))
+			if index < 0 || index >= l {
+				return nil, fmt.Errorf("index out of range for %v: %d", v, index)
+			}
+			return v[index], nil
+		default:
+			// Otherwise, we index into the string
+			str := toString(v)
+			l := int64(len(str))
+			if index < 0 || index >= l {
+				return nil, fmt.Errorf("index out of range for %v: %d", v, index)
+			}
+			return string(str[index]), nil
+		}
+	}))
+
 	MustAddFunction(NewWithPositionalArgs("len", 1, func(args ...interface{}) (interface{}, error) {
-		length := len(toString(args[0]))
+		var length int
+		value := reflect.ValueOf(args[0])
+		switch value.Kind() {
+		case reflect.Slice:
+			length = value.Len()
+		case reflect.Map:
+			length = value.Len()
+		default:
+			length = len(toString(args[0]))
+		}
 		return float64(length), nil
 	}))
 

diff --git a/dsl_test.go b/dsl_test.go
@@ -11,6 +11,15 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func TestIndex(t *testing.T) {
+	index, err := govaluate.NewEvaluableExpressionWithFunctions("index(split(url, '.', -1), 1) == 'example'", DefaultHelperFunctions)
+	require.Nil(t, err, "could not compile index")
+
+	result, err := index.Evaluate(map[string]interface{}{"url": "https://www.example.com"})
+	require.Nil(t, err, "could not evaluate index")
+	require.Equal(t, true, result, "could not get index data")
+}
+
 func TestDSLURLEncodeDecode(t *testing.T) {
 	encoded, err := DefaultHelperFunctions["url_encode"]("&test\"")
 	require.Nil(t, err, "could not url encode")
@@ -113,6 +122,7 @@ func TestGetPrintableDslFunctionSignatures(t *testing.T) {
 	hmac(arg1, arg2, arg3 interface{}) interface{}
 	html_escape(arg1 interface{}) interface{}
 	html_unescape(arg1 interface{}) interface{}
+	index(arg1, arg2 interface{}) interface{}
 	ip_format(arg1, arg2 interface{}) interface{}
 	join(separator string, elements ...interface{}) string
 	join(separator string, elements []interface{}) string
@@ -227,6 +237,7 @@ func TestDslExpressions(t *testing.T) {
 		`hex_decode("6161")`:                                      "aa",
 		`len("Hello")`:                                            float64(5),
 		`len(1234)`:                                               float64(4),
+		`len(split("1.2.3.4",'.',-1))`:                            float64(4),
 		`contains("Hello", "lo")`:                                 true,
 		`starts_with("Hello", "He")`:                              true,
 		`ends_with("Hello", "lo")`:                                true,