Add substr function (Velocidex#1059)

vql/functions/strings.go

@@ -56,6 +56,60 @@ func (self StripFunction) Info(scope vfilter.Scope, type_map *vfilter.TypeMap) *
 	}
 }
 
+type SubStrFunction struct{}
+
+type SubStrArgs struct {
+	Str   string `vfilter:"required,field=str,doc=The string to shorten"`
+	Start int    `vfilter:"optional,field=start,doc=Beginning index of substring"`
+	End   int    `vfilter:"optional,field=end,doc=End index of substring"`
+}
+
+func (self *SubStrFunction) Call(ctx context.Context,
+	scope vfilter.Scope,
+	args *ordereddict.Dict) vfilter.Any {
+	arg := &SubStrArgs{}
+	err := vfilter.ExtractArgs(scope, args, arg)
+	if err != nil {
+		scope.Log("substr: %s", err.Error())
+		return nil
+	}
+
+	if arg.Start < 0 || arg.End < 0 {
+		scope.Log("substr: Start and End values must be greater than 0")
+		return nil
+	}
+
+	if arg.Start == 0 && arg.End == 0 {
+		return arg.Str
+	} else if arg.Start != 0 && arg.End == 0 {
+		arg.End = len(arg.Str)
+	} else if arg.End < arg.Start {
+		scope.Log("substr: End must be greater than start!")
+		return nil
+	}
+
+	counter, start_index := 0, 0
+	for i := range arg.Str {
+		if counter == arg.Start {
+			start_index = i
+		}
+		if counter == arg.End {
+			return arg.Str[start_index:i]
+		}
+		counter++
+	}
+	return arg.Str[start_index:]
+}
+
+func (self *SubStrFunction) Info(scope vfilter.Scope, type_map *vfilter.TypeMap) *vfilter.FunctionInfo {
+	return &vfilter.FunctionInfo{
+		Name:    "substr",
+		Doc:     "Create a substring from a string",
+		ArgType: type_map.AddType(scope, &SubStrArgs{}),
+	}
+}
+
 func init() {
 	vql_subsystem.RegisterFunction(&StripFunction{})
+	vql_subsystem.RegisterFunction(&SubStrFunction{})
 }