Refactor and improve on VQL reformatter.

fixtures/vql_queries.golden

@@ -323,7 +323,7 @@
       "bar": 0
     }
   ],
-  "035 Limit and order: SELECT * FROM test() ORDER BY foo DESC  LIMIT 1 ": [
+  "035 Limit and order: SELECT * FROM test() ORDER BY foo DESC LIMIT 1 ": [
     {
       "foo": 4,
       "bar": 2

protocols/protocol_membership.go

@@ -4,6 +4,7 @@ import (
 	"reflect"
 	"strings"
 
+	"github.com/Velocidex/ordereddict"
 	"www.velocidex.com/golang/vfilter/types"
 )
 
@@ -30,6 +31,13 @@ func (self MembershipDispatcher) Membership(scope types.Scope, a types.Any, b ty
 	case types.Null, *types.Null, nil:
 		return false
 
+	case *ordereddict.Dict:
+		a_str, ok := a.(string)
+		if ok {
+			_, pres := t.Get(a_str)
+			return pres
+		}
+
 	case string:
 		// 'he' in 'hello'
 		a_str, ok := a.(string)

reformat/fixtures/formatting.golden

@@ -8,6 +8,7 @@ FROM ALLUploads
 
 
 1 MultiLine Comment:
+
 /*
    This is a multiline comment
 */
@@ -33,10 +34,8 @@ LET LocalDeviceMajor <= (253, 7, -- loop
 
 4 Comment above plugin arg:
 SELECT *
-FROM info(
-  -- This is a comment for Foo
-  Foo=1,
-  Bar=2)
+FROM info(-- This is a comment for Foo
+Foo=1, Bar=2)
 
 
 
@@ -85,18 +84,18 @@ FROM Artifact.Windows.EventLogs.EvtxHunter(
 
 
 11 If Plugin:
-LET DateAfterTime <= if(
-    condition=TargetTime,
-    then=timestamp(epoch=TargetTime.Unix - TargetTimeBox),
-    else=timestamp(epoch=now() - TargetTimeBox))
+LET DateAfterTime <= if(condition=TargetTime,
+                        then=timestamp(epoch=TargetTime.Unix - TargetTimeBox),
+                        else=timestamp(epoch=now() - TargetTimeBox))
 
 
 
 12 If Plugin with longer args:
-LET DateAfterTime <= if(
-    condition=TargetTime,
-    then=timestamp(epoch=TargetTime.Unix - TargetTimeBox + SomeConstant),
-    else=timestamp(epoch=now() - TargetTimeBox))
+LET DateAfterTime <= if(condition=TargetTime,
+                        then=timestamp(epoch=TargetTime.Unix -
+                                         TargetTimeBox +
+                                         SomeConstant),
+                        else=timestamp(epoch=now() - TargetTimeBox))
 
 
 
@@ -155,32 +154,22 @@ LET Foo(x, y, z) = SELECT *
 
 20 Nesting:
 SELECT *
-FROM foreach(
-  row={
+FROM foreach(row={
     SELECT *
-    FROM foreach(
-      row={
-        SELECT *
-        FROM info
-      },
-      query={
-        SELECT *
-        FROM info()
-      })
+    FROM foreach(row={ SELECT * FROM info }, query={ SELECT * FROM info() })
   },
-  query=ForeachQuery(X=1))
+             query=ForeachQuery(X=1))
 
 
 
 21 Foreach:
 SELECT *
-FROM foreach(
-  row={
+FROM foreach(row={
     SELECT *
     FROM clients()
-    LIMIT 5 
+    LIMIT 5
   },
-  query={
+             query={
     SELECT client_info(client_id=client_id)
     FROM glob(globs=AGlob)
   })
@@ -189,9 +178,9 @@ FROM foreach(
 
 22 Subquery:
 SELECT {
-         SELECT *
-         FROM info()
-       } AS Foo,
+    SELECT *
+    FROM info()
+    } AS Foo,
        Bar
 FROM scope()
 
@@ -205,7 +194,7 @@ SELECT A AS First,
 FROM info(arg=1, arg2=3)
 WHERE 1
 ORDER BY C
-LIMIT 1 
+LIMIT 1
 
 
 
@@ -215,3 +204,28 @@ FROM scope()
 
 
 
+25 Very long LET with SELECT pushing far to right:
+LET this_is_a_long_name = SELECT
+    *,
+    upload_directory(
+      accessor="collector",
+      file=RootPathSpec + "hello world " + _Components) AS UploadedFile
+  FROM ALLUploads
+
+
+
+26 Complex long lines. Args should line up on the function they are in.:
+LET enumerate_path = SELECT
+    regex_replace(source=TargetPath,
+                  re='''\%USERPROFILE\%''',
+                  replace=Directory) AS TargetPath,
+    *,
+    check_exist(path=regex_replace(source=TargetPath,
+                                   re='''\%USERPROFILE\%''',
+                                   replace=Directory))[0] AS Exists,
+    MaxSize - rand(range=(MaxSize - MinSize)) - len(
+      list=unhex(string=MagicBytes)) - 7 AS _PaddingSize
+  FROM Honeyfiles
+
+
+

reformat/visitor.go

@@ -1,6 +1,8 @@
 package reformat
 
 import (
+	"strings"
+
 	"www.velocidex.com/golang/vfilter"
 	"www.velocidex.com/golang/vfilter/types"
 )
@@ -15,5 +17,11 @@ func ReFormatVQL(scope types.Scope, query string,
 	visitor := vfilter.NewVisitor(scope, options)
 	visitor.Visit(vql)
 
-	return visitor.ToString(), nil
+	lines := strings.Split(visitor.ToString(), "\n")
+	result := make([]string, 0, len(lines))
+	for _, l := range lines {
+		result = append(result, strings.TrimRight(l, " "))
+	}
+
+	return strings.Join(result, "\n"), nil
 }

reformat/visitor_test.go

@@ -85,6 +85,22 @@ FROM Artifact.Windows.EventLogs.EvtxHunter(EvtxGlob='''%SystemRoot%\System32\Win
 	{"Subquery", "SELECT {SELECT * FROM info()} AS Foo, Bar FROM scope()"},
 	{"Simple Statement", "SELECT A AS First, B AS Second, C, D FROM info(arg=1, arg2=3) WHERE 1 ORDER BY C LIMIT 1"},
 	{"Explain statements", "EXPLAIN SELECT 'A' FROM scope()"},
+	{"Very long LET with SELECT pushing far to right", `
+LET this_is_a_long_name = SELECT *, upload_directory(accessor="collector", file=RootPathSpec + "hello world " + _Components) AS UploadedFile FROM ALLUploads
+`},
+	{"Complex long lines. Args should line up on the function they are in.", `
+ LET enumerate_path = SELECT
+           regex_replace(source=TargetPath,
+                         re='''\%USERPROFILE\%''',
+                         replace=Directory) AS TargetPath,
+           *,
+           check_exist(path=regex_replace(source=TargetPath,
+                                          re='''\%USERPROFILE\%''',
+                                          replace=Directory))[0] AS Exists,
+           MaxSize - rand(range=(MaxSize - MinSize)) - len(
+             list=unhex(string=MagicBytes)) - 7 AS _PaddingSize
+         FROM Honeyfiles
+`},
 }
 
 func makeTestScope() types.Scope {
@@ -100,7 +116,7 @@ func TestVQLQueries(t *testing.T) {
 	golden := ""
 
 	for idx, testCase := range reformatTests {
-		if false && idx != 24 {
+		if false && idx != 20 {
 			continue
 		}
 

utils/debug.go

@@ -1,3 +1,15 @@
 package utils
 
+import "fmt"
+
+const (
+	debug = false
+)
+
 func DlvBreak() {}
+
+func DebugPrint(format string) {
+	if debug {
+		fmt.Println(format)
+	}
+}

utils/dict/dict.go

@@ -89,13 +89,18 @@ func normalize_value(ctx context.Context,
 		return result
 
 	default:
-		a_value := reflect.Indirect(reflect.ValueOf(value))
-		a_type := a_value.Type()
+		a_value := reflect.ValueOf(value)
+		a_value = reflect.Indirect(a_value)
+		a_type := reflect.TypeOf(value)
 		if a_type == nil {
 			return types.Null{}
 		}
 
 		if a_type.Kind() == reflect.Slice || a_type.Kind() == reflect.Array {
+			if types.IsNil(a_value) {
+				return types.Null{}
+			}
+
 			length := a_value.Len()
 			result := make([]types.Any, 0, length)
 			for i := 0; i < length; i++ {
@@ -106,6 +111,10 @@ func normalize_value(ctx context.Context,
 
 			// Map keys are random so they lead to unstable output
 		} else if a_type.Kind() == reflect.Map {
+			if types.IsNil(a_value) {
+				return types.Null{}
+			}
+
 			result := ordereddict.NewDict()
 			for _, key := range a_value.MapKeys() {
 				str_key, ok := key.Interface().(string)

vfilter_test.go

@@ -245,6 +245,8 @@ var execTestsSerialization = []execTest{
 		Set("foo", int64(1)).
 		Set("bar", []Any{int64(2), int64(3)})},
 
+	{"'foo' IN dict(foo=0)", true},
+
 	// Associative
 	// Relies on pre-populating the scope with a Dict.
 	{"foo.bar.baz, foo.bar2", []float64{5, 7}},