feat(vql): add reformat function

artifacts/testdata/server/testcases/reformat.in.yaml

@@ -0,0 +1,63 @@
+# This file contains the same tests that are in services/repository/reformat_test.go
+
+Parameters:
+  Simple: |
+    name: Foo
+    sources:
+    - query: |
+        SELECT A,B,C
+        FROM info(arg1=Foobar, arg2="XXXXX")
+    - query: |-
+        SELECT A,B,C      FROM info()
+    - description: Foo bar
+
+  SingleLine: |
+    name: Single line
+    sources:
+    - query: SELECT * FROM info()
+    - query: |
+        SELECT A FROM scope()
+
+  Exported: |
+    name: Exported
+    export: |
+      SELECT * FROM scope()
+
+  Preconditions: |
+    name: Preconditions
+    precondition: SELECT * FROM info()
+    sources:
+    - precondition: |
+        SELECT * FROM info()
+
+  Notebook: |
+    name: Notebook
+    sources:
+    - query: |
+        SELECT A,B,C
+        FROM scope()
+      notebook:
+        - name: Test
+          type: vql_suggestion
+          template: |
+            SELECT * FROM scope()
+
+  ColumnTypes: |
+    name: Column Types
+    sources:
+    - query: |
+        SELECT A,B,C
+        FROM scope()
+    column_types:
+      - name: A
+        type: string
+
+Queries:
+  - SELECT 
+      split(string=reformat(artifact=Simple).Artifact, sep="\\n") AS Simple,
+      split(string=reformat(artifact=SingleLine).Artifact, sep="\\n") AS SingleLine,
+      split(string=reformat(artifact=Exported).Artifact, sep="\\n") AS Exported,
+      split(string=reformat(artifact=Preconditions).Artifact, sep="\\n") AS Preconditions,
+      split(string=reformat(artifact=Notebook).Artifact, sep="\\n") AS Notebook,
+      split(string=reformat(artifact=ColumnTypes).Artifact, sep="\\n") AS ColumnTypes
+    FROM scope()

artifacts/testdata/server/testcases/reformat.out.yaml

@@ -0,0 +1,69 @@
+Query: SELECT split(string=reformat(artifact=Simple).Artifact, sep="\\n") AS Simple, split(string=reformat(artifact=SingleLine).Artifact, sep="\\n") AS SingleLine, split(string=reformat(artifact=Exported).Artifact, sep="\\n") AS Exported, split(string=reformat(artifact=Preconditions).Artifact, sep="\\n") AS Preconditions, split(string=reformat(artifact=Notebook).Artifact, sep="\\n") AS Notebook, split(string=reformat(artifact=ColumnTypes).Artifact, sep="\\n") AS ColumnTypes FROM scope()
+Output: [
+ {
+  "Simple": [
+   "name: Foo",
+   "sources:",
+   "- query: |",
+   "    SELECT A,",
+   "           B,",
+   "           C",
+   "    FROM info(arg1=Foobar, arg2=\"XXXXX\")",
+   "- query: |-",
+   "    SELECT A,",
+   "           B,",
+   "           C",
+   "    FROM info()",
+   "- description: Foo bar"
+  ],
+  "SingleLine": [
+   "name: Single line",
+   "sources:",
+   "- query: SELECT * FROM info()",
+   "- query: |",
+   "    SELECT A",
+   "    FROM scope()"
+  ],
+  "Exported": [
+   "name: Exported",
+   "export: |",
+   "  SELECT *",
+   "  FROM scope()"
+  ],
+  "Preconditions": [
+   "name: Preconditions",
+   "precondition: SELECT * FROM info()",
+   "sources:",
+   "- precondition: |",
+   "    SELECT *",
+   "    FROM info()"
+  ],
+  "Notebook": [
+   "name: Notebook",
+   "sources:",
+   "- query: |",
+   "    SELECT A,",
+   "           B,",
+   "           C",
+   "    FROM scope()",
+   "  notebook:",
+   "    - name: Test",
+   "      type: vql_suggestion",
+   "      template: |",
+   "        SELECT * FROM scope()"
+  ],
+  "ColumnTypes": [
+   "name: Column Types",
+   "sources:",
+   "- query: |",
+   "    SELECT A,",
+   "           B,",
+   "           C",
+   "    FROM scope()",
+   "column_types:",
+   "  - name: A",
+   "    type: string"
+  ]
+ }
+]
+

bin/reformat.go

@@ -2,19 +2,21 @@ package main
 
 import (
 	"fmt"
-	"os"
+	"log"
+	"path/filepath"
 
+	"github.com/Velocidex/ordereddict"
 	"www.velocidex.com/golang/velociraptor/config"
-	"www.velocidex.com/golang/velociraptor/constants"
 	logging "www.velocidex.com/golang/velociraptor/logging"
 	"www.velocidex.com/golang/velociraptor/services"
 	"www.velocidex.com/golang/velociraptor/startup"
-	"www.velocidex.com/golang/velociraptor/utils"
+	"www.velocidex.com/golang/velociraptor/vql/acl_managers"
 )
 
 var (
-	reformat      = artifact_command.Command("reformat", "Reformat a set of artifacts")
-	reformat_args = reformat.Arg("paths", "Paths to artifact yaml files").Required().Strings()
+	reformat         = artifact_command.Command("reformat", "Reformat a set of artifacts")
+	reformat_dry_run = reformat.Flag("dry", "Do not overwrite files, just report errors").Bool()
+	reformat_args    = reformat.Arg("paths", "Paths to artifact yaml files").Required().Strings()
 )
 
 func doReformat() error {
@@ -38,60 +40,47 @@ func doReformat() error {
 		return err
 	}
 
-	manager, err := services.GetRepositoryManager(config_obj)
-	if err != nil {
-		return err
-	}
-
 	logger := logging.GetLogger(config_obj, &logging.ToolComponent)
 
-	// Report all errors and keep going as much as possible.
-	returned_errs := make(map[string]error)
-
+	var artifact_paths []string
 	for _, artifact_path := range *reformat_args {
-		returned_errs[artifact_path] = nil
-
-		fd, err := os.Open(artifact_path)
-		if err != nil {
-			returned_errs[artifact_path] = err
-			continue
-		}
-
-		data, err := utils.ReadAllWithLimit(fd, constants.MAX_MEMORY)
+		abs, err := filepath.Abs(artifact_path)
 		if err != nil {
-			returned_errs[artifact_path] = err
-			fd.Close()
+			logger.Error("reformat: could not get absolute path for %v", artifact_path)
 			continue
 		}
-		fd.Close()
 
-		reformatted, err := manager.ReformatVQL(ctx, string(data))
-		if err != nil {
-			returned_errs[artifact_path] = err
-			continue
-		}
+		artifact_paths = append(artifact_paths, abs)
+	}
 
-		out_fd, err := os.OpenFile(artifact_path,
-			os.O_TRUNC|os.O_WRONLY|os.O_CREATE, 0600)
-		if err != nil {
-			returned_errs[artifact_path] = err
-			continue
-		}
-		_, _ = out_fd.Write([]byte(reformatted))
-		out_fd.Close()
+	artifact_logger := &LogWriter{config_obj: sm.Config}
+	builder := services.ScopeBuilder{
+		Config:     sm.Config,
+		ACLManager: acl_managers.NewRoleACLManager(sm.Config, "administrator"),
+		Logger:     log.New(artifact_logger, "", 0),
+		Env: ordereddict.NewDict().
+			Set("DryRun", *reformat_dry_run).
+			Set("Artifacts", artifact_paths),
 	}
 
-	var ret error
-	for artifact_path, err := range returned_errs {
-		if err != nil {
-			logger.Error("%v: <red>%v</>", artifact_path, err)
-			ret = err
-		} else {
-			logger.Info("Reformatted %v: <green>OK</>", artifact_path)
-		}
+	query := `
+      SELECT reformat(artifact=read_file(filename=_value)) AS Result
+        FROM foreach(row=Artifacts)
+      WHERE if(condition=Result.Error,
+          then=log(level="ERROR", message="%v: <red>%v</>",
+                   args=[_value, Result.Error], dedup=-1),
+          else=log(message="Reformatted %v: <green>OK</>",
+                   args=_value, dedup=-1)
+               AND NOT DryRun
+               AND copy(accessor="data", dest=_value, filename=Result.Artifact))
+       AND FALSE
+    `
+	err = runQueryWithEnv(query, builder, "json")
+	if err != nil {
+		logger.Error("reformat: error running query: %v", query)
 	}
 
-	return ret
+	return artifact_logger.Error
 }
 
 func init() {

docs/references/vql.yaml

@@ -8621,6 +8621,23 @@
   - linux_amd64_cgo
   - windows_386_cgo
   - windows_amd64_cgo
+- name: reformat
+  description: |-
+    Reformat VQL
+
+    This function will reformat the artifact provided and return the reformatted content.
+  type: Function
+  args:
+  - name: artifact
+    type: string
+    description: The artifact VQL to reformat.
+    required: true
+  platforms:
+  - darwin_amd64_cgo
+  - darwin_arm64_cgo
+  - linux_amd64_cgo
+  - windows_386_cgo
+  - windows_amd64_cgo
 - name: reg
   description: |
     An alias for the `registry` accessor, which accesses the registry using the

vql/functions/reformat.go

@@ -0,0 +1,82 @@
+package functions
+
+import (
+	"context"
+	"strings"
+
+	"github.com/Velocidex/ordereddict"
+	"www.velocidex.com/golang/velociraptor/services"
+	vql_subsystem "www.velocidex.com/golang/velociraptor/vql"
+	"www.velocidex.com/golang/vfilter"
+	"www.velocidex.com/golang/vfilter/arg_parser"
+)
+
+type ReformatFunctionArgs struct {
+	Artifact string `vfilter:"required,field=artifact,doc=The artifact VQL to reformat."`
+}
+
+type ReformatFunction struct{}
+
+type ReformatFunctionResult struct {
+	Artifact string
+	Error    string
+}
+
+func (self *ReformatFunctionResult) ToDict() *ordereddict.Dict {
+	return ordereddict.NewDict().
+		Set("Artifact", self.Artifact).
+		Set("Error", self.Error)
+}
+
+func (self *ReformatFunction) Call(ctx context.Context, scope vfilter.Scope, args *ordereddict.Dict) vfilter.Any {
+	defer vql_subsystem.RegisterMonitor(ctx, "reformat", args)()
+
+	result := &ReformatFunctionResult{}
+
+	arg := &ReformatFunctionArgs{}
+	err := arg_parser.ExtractArgsWithContext(ctx, scope, args, arg)
+	if err != nil {
+		result.Artifact = arg.Artifact
+		result.Error = err.Error()
+		return result.ToDict()
+	}
+
+	config_obj, ok := vql_subsystem.GetServerConfig(scope)
+	if !ok {
+		scope.Log("reformat: Must be run on the server")
+		return vfilter.Null{}
+	}
+
+	manager, err := services.GetRepositoryManager(config_obj)
+	if err != nil {
+		result.Artifact = arg.Artifact
+		result.Error = err.Error()
+		return result.ToDict()
+	}
+
+	reformatted, err := manager.ReformatVQL(ctx, arg.Artifact)
+	if err != nil {
+		result.Artifact = arg.Artifact
+		result.Error = err.Error()
+		return result.ToDict()
+	}
+
+	result.Artifact = strings.Trim(reformatted, "\n")
+	result.Error = ""
+
+	return result.ToDict()
+}
+
+func (self ReformatFunction) Info(scope vfilter.Scope, type_map *vfilter.TypeMap) *vfilter.FunctionInfo {
+	return &vfilter.FunctionInfo{
+		Name: "reformat",
+		Doc: `Reformat VQL
+
+This function will reformat the artifact provided and return the reformatted content.`,
+		ArgType: type_map.AddType(scope, &ReformatFunctionArgs{}),
+	}
+}
+
+func init() {
+	vql_subsystem.RegisterFunction(&ReformatFunction{})
+}