Adding python Unpack implementation

dsl.go

@@ -28,6 +28,7 @@ import (
 	"time"
 
 	"github.com/Knetic/govaluate"
+	"github.com/Mzack9999/gostruct"
 	"github.com/asaskevich/govalidator"
 	"github.com/hashicorp/go-version"
 	"github.com/kataras/jwt"
@@ -772,10 +773,16 @@ func init() {
 				return nil, ErrInvalidDslFunction
 			}
 			argStr := toString(args[0])
+			if len(argStr) == 0 {
+				return nil, errors.New("empty string")
+			}
 			start, err := strconv.Atoi(toString(args[1]))
 			if err != nil {
 				return nil, errors.Wrap(err, "invalid start position")
 			}
+			if start > len(argStr) {
+				return nil, errors.Wrap(err, "start position bigger than slice length")
+			}
 			if len(args) == 2 {
 				return argStr[start:], nil
 			}
@@ -785,7 +792,13 @@ func init() {
 				return nil, errors.Wrap(err, "invalid end position")
 			}
 			if end < 0 {
-				end += len(argStr)
+				return nil, errors.Wrap(err, "negative end position")
+			}
+			if end < start {
+				return nil, errors.Wrap(err, "end position before start")
+			}
+			if end > len(argStr) {
+				return nil, errors.Wrap(err, "end position bigger than slice length start")
 			}
 			return argStr[start:end], nil
 		}))
@@ -947,9 +960,73 @@ func init() {
 		return formattedIps[0], nil
 	}))
 	MustAddFunction(NewWithPositionalArgs("llm_prompt", 1, func(args ...interface{}) (interface{}, error) {
-		prompt := args[0].(string)
+		prompt, ok := args[0].(string)
+		if !ok {
+			return nil, errors.New("invalid prompt")
+		}
 		return llm.Query(prompt)
 	}))
+	MustAddFunction(NewWithPositionalArgs("unpack", 2, func(args ...interface{}) (interface{}, error) {
+		// format as string (ref: https://docs.python.org/3/library/struct.html#format-characters)
+		format, ok := args[0].(string)
+		if !ok {
+			return nil, errors.New("invalid format")
+		}
+		// binary packed data
+		data, ok := args[1].(string)
+		if !ok {
+			return nil, errors.New("invalid data")
+		}
+		// convert flat format into slice (eg. ">I" => [">","I"])
+		var formatParts []string
+		for idx := range format {
+			formatParts = append(formatParts, string(format[idx]))
+		}
+		// the dsl function supports unpacking only one type at a time
+		unpackedData, err := gostruct.UnPack(formatParts, []byte(data))
+		if len(unpackedData) > 0 {
+			return unpackedData[0], err
+		}
+		return nil, errors.New("no result")
+	}))
+	MustAddFunction(NewWithSingleSignature("xor",
+		"(args ...interface{}) interface{}",
+		func(args ...interface{}) (interface{}, error) {
+			if len(args) < 2 {
+				return nil, errors.New("at least two arguments needed")
+			}
+
+			n := -1
+			for _, arg := range args {
+				var b []byte
+				switch v := arg.(type) {
+				case string:
+					b = []byte(v)
+				case []byte:
+					b = v
+				default:
+					return nil, fmt.Errorf("invalid argument type %T", arg)
+				}
+				if n == -1 {
+					n = len(b)
+				} else if len(b) != n {
+					return nil, errors.New("all arguments must have the same length")
+				}
+			}
+
+			result := make([]byte, n)
+			for i := 0; i < n; i++ {
+				for _, arg := range args {
+					b, ok := arg.([]byte)
+					if !ok {
+						b = []byte(arg.(string))
+					}
+					result[i] ^= b[i]
+				}
+			}
+
+			return result, nil
+		}))
 
 	DefaultHelperFunctions = HelperFunctions()
 	FunctionNames = GetFunctionNames(DefaultHelperFunctions)

dsl_test.go

@@ -166,9 +166,11 @@ func TestGetPrintableDslFunctionSignatures(t *testing.T) {
 	uniq(input number) string
 	uniq(input string) string
 	unix_time(optionalSeconds uint) float64
+	unpack(arg1, arg2 interface{}) interface{}
 	url_decode(arg1 interface{}) interface{}
 	url_encode(arg1 interface{}) interface{}
 	wait_for(seconds uint)
+	xor(args ...interface{}) interface{}
 	zlib(arg1 interface{}) interface{}
 	zlib_decode(arg1 interface{}) interface{}
 `
@@ -251,7 +253,7 @@ func TestDslExpressions(t *testing.T) {
 		`hmac('sha256', 'test', 'scrt')`:                          "1f1bff5574f18426eb376d6dd5368a754e67a798aa2074644d5e3fd4c90c7a92",
 		`hmac('sha512', 'test', 'scrt')`:                          "1d3fff1dbb7369c1615ffb494813146bea051ce07e5d44bdeca539653ea97656bf9d38db264cddbe6a83ea15139c8f861a7e73e10e43ad4865e852a9ee6de2e9",
 		`substr('xxtestxxx',2)`:                                   "testxxx",
-		`substr('xxtestxxx',2,-2)`:                                "testx",
+		`substr('xxtestxxx',2,4)`:                                 "te",
 		`substr('xxtestxxx',2,6)`:                                 "test",
 		`sort(12453)`:                                             "12345",
 		`sort("a1b2c3d4e5")`:                                      "12345abcde",
@@ -274,6 +276,8 @@ func TestDslExpressions(t *testing.T) {
 		`ip_format('127.0.0.1', '3')`:                                          "0177.0.0.01",
 		`ip_format('127.0.0.1', '5')`:                                          "281472812449793",
 		`ip_format('127.0.1.0', '11')`:                                         "127.0.256",
+		"unpack('>I', '\xac\xd7\t\xd0')":                                       -272646673,
+		"xor('\x01\x02', '\x02\x01')":                                          []uint8([]byte{0x3, 0x3}),
 	}
 
 	testDslExpressions(t, dslExpressions)

go.mod

@@ -4,6 +4,7 @@ go 1.19
 
 require (
 	github.com/Knetic/govaluate v3.0.0+incompatible
+	github.com/Mzack9999/gostruct v0.0.0-20230415193108-30b70932da81
 	github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2
 	github.com/hashicorp/go-version v1.6.0
 	github.com/kataras/jwt v0.1.8

go.sum

@@ -1,5 +1,7 @@
 github.com/Knetic/govaluate v3.0.0+incompatible h1:7o6+MAPhYTCF0+fdvoz1xDedhRb4f6s9Tn1Tt7/WTEg=
 github.com/Knetic/govaluate v3.0.0+incompatible/go.mod h1:r7JcOSlj0wfOMncg0iLm8Leh48TZaKVeNIfJntJ2wa0=
+github.com/Mzack9999/gostruct v0.0.0-20230415193108-30b70932da81 h1:rwHZjxG8Cx3+FNujiZRuJbYTLHmW8U9+6xIoTseKA/I=
+github.com/Mzack9999/gostruct v0.0.0-20230415193108-30b70932da81/go.mod h1:iXPMmoXMc0ZsSmbbHqhWCWd8w7FkXM7DU2IBf5OS+5g=
 github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2 h1:DklsrG3dyBCFEj5IhUbnKptjxatkF07cF2ak3yi77so=
 github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2/go.mod h1:WaHUgvxTVq04UNunO+XhnAqY/wQc+bxr74GqbsZ/Jqw=
 github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuPk=