Skip to content

Commit d7c0225

Browse files
nullifysecurityscudette
nullifysecurity
and
scudette
authored
feat(vql): add reformat function (#4675)
This change introduces a new `reformat` function that wraps the existing VQL reformatting logic, allowing artifact authors to call it directly from queries to reformat their artifacts. It also updates the `artifacts reformat` command to align with Velociraptor’s general approach, keeping the CLI as a thin layer over the underlying VQL functionality. Right now, the `reformat` function accepts the artifact content as a string, similar to how the `verify` function works. I’m wondering whether it would make more sense to switch this to a `file` argument instead and let the caller optionally provide an `accessor`. Would that be a better fit and provide more flexibility? As always, let me know if there is anything you would like me to clarify or change! 😃 --------- Co-authored-by: Mike Cohen <mike@velocidex.com>
1 parent 9ddb189 commit d7c0225

File tree

5 files changed

+267
-47
lines changed

5 files changed

+267
-47
lines changed

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

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
# This file contains the same tests that are in services/repository/reformat_test.go
2+
3+
Parameters:
4+
Simple: |
5+
name: Foo
6+
sources:
7+
- query: |
8+
SELECT A,B,C
9+
FROM info(arg1=Foobar, arg2="XXXXX")
10+
- query: |-
11+
SELECT A,B,C FROM info()
12+
- description: Foo bar
13+
14+
SingleLine: |
15+
name: Single line
16+
sources:
17+
- query: SELECT * FROM info()
18+
- query: |
19+
SELECT A FROM scope()
20+
21+
Exported: |
22+
name: Exported
23+
export: |
24+
SELECT * FROM scope()
25+
26+
Preconditions: |
27+
name: Preconditions
28+
precondition: SELECT * FROM info()
29+
sources:
30+
- precondition: |
31+
SELECT * FROM info()
32+
33+
Notebook: |
34+
name: Notebook
35+
sources:
36+
- query: |
37+
SELECT A,B,C
38+
FROM scope()
39+
notebook:
40+
- name: Test
41+
type: vql_suggestion
42+
template: |
43+
SELECT * FROM scope()
44+
45+
ColumnTypes: |
46+
name: Column Types
47+
sources:
48+
- query: |
49+
SELECT A,B,C
50+
FROM scope()
51+
column_types:
52+
- name: A
53+
type: string
54+
55+
Queries:
56+
- SELECT
57+
split(string=reformat(artifact=Simple).Artifact, sep="\\n") AS Simple,
58+
split(string=reformat(artifact=SingleLine).Artifact, sep="\\n") AS SingleLine,
59+
split(string=reformat(artifact=Exported).Artifact, sep="\\n") AS Exported,
60+
split(string=reformat(artifact=Preconditions).Artifact, sep="\\n") AS Preconditions,
61+
split(string=reformat(artifact=Notebook).Artifact, sep="\\n") AS Notebook,
62+
split(string=reformat(artifact=ColumnTypes).Artifact, sep="\\n") AS ColumnTypes
63+
FROM scope()

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

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
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()
2+
Output: [
3+
{
4+
"Simple": [
5+
"name: Foo",
6+
"sources:",
7+
"- query: |",
8+
" SELECT A,",
9+
" B,",
10+
" C",
11+
" FROM info(arg1=Foobar, arg2=\"XXXXX\")",
12+
"- query: |-",
13+
" SELECT A,",
14+
" B,",
15+
" C",
16+
" FROM info()",
17+
"- description: Foo bar"
18+
],
19+
"SingleLine": [
20+
"name: Single line",
21+
"sources:",
22+
"- query: SELECT * FROM info()",
23+
"- query: |",
24+
" SELECT A",
25+
" FROM scope()"
26+
],
27+
"Exported": [
28+
"name: Exported",
29+
"export: |",
30+
" SELECT *",
31+
" FROM scope()"
32+
],
33+
"Preconditions": [
34+
"name: Preconditions",
35+
"precondition: SELECT * FROM info()",
36+
"sources:",
37+
"- precondition: |",
38+
" SELECT *",
39+
" FROM info()"
40+
],
41+
"Notebook": [
42+
"name: Notebook",
43+
"sources:",
44+
"- query: |",
45+
" SELECT A,",
46+
" B,",
47+
" C",
48+
" FROM scope()",
49+
" notebook:",
50+
" - name: Test",
51+
" type: vql_suggestion",
52+
" template: |",
53+
" SELECT * FROM scope()"
54+
],
55+
"ColumnTypes": [
56+
"name: Column Types",
57+
"sources:",
58+
"- query: |",
59+
" SELECT A,",
60+
" B,",
61+
" C",
62+
" FROM scope()",
63+
"column_types:",
64+
" - name: A",
65+
" type: string"
66+
]
67+
}
68+
]
69+

‎bin/reformat.go‎

Lines changed: 36 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -2,19 +2,21 @@ package main
22

33
import (
44
"fmt"
5-
"os"
5+
"log"
6+
"path/filepath"
67

8+
"github.com/Velocidex/ordereddict"
79
"www.velocidex.com/golang/velociraptor/config"
8-
"www.velocidex.com/golang/velociraptor/constants"
910
logging "www.velocidex.com/golang/velociraptor/logging"
1011
"www.velocidex.com/golang/velociraptor/services"
1112
"www.velocidex.com/golang/velociraptor/startup"
12-
"www.velocidex.com/golang/velociraptor/utils"
13+
"www.velocidex.com/golang/velociraptor/vql/acl_managers"
1314
)
1415

1516
var (
16-
reformat = artifact_command.Command("reformat", "Reformat a set of artifacts")
17-
reformat_args = reformat.Arg("paths", "Paths to artifact yaml files").Required().Strings()
17+
reformat = artifact_command.Command("reformat", "Reformat a set of artifacts")
18+
reformat_dry_run = reformat.Flag("dry", "Do not overwrite files, just report errors").Bool()
19+
reformat_args = reformat.Arg("paths", "Paths to artifact yaml files").Required().Strings()
1820
)
1921

2022
func doReformat() error {
@@ -38,60 +40,47 @@ func doReformat() error {
3840
return err
3941
}
4042

41-
manager, err := services.GetRepositoryManager(config_obj)
42-
if err != nil {
43-
return err
44-
}
45-
4643
logger := logging.GetLogger(config_obj, &logging.ToolComponent)
4744

48-
// Report all errors and keep going as much as possible.
49-
returned_errs := make(map[string]error)
50-
45+
var artifact_paths []string
5146
for _, artifact_path := range *reformat_args {
52-
returned_errs[artifact_path] = nil
53-
54-
fd, err := os.Open(artifact_path)
55-
if err != nil {
56-
returned_errs[artifact_path] = err
57-
continue
58-
}
59-
60-
data, err := utils.ReadAllWithLimit(fd, constants.MAX_MEMORY)
47+
abs, err := filepath.Abs(artifact_path)
6148
if err != nil {
62-
returned_errs[artifact_path] = err
63-
fd.Close()
49+
logger.Error("reformat: could not get absolute path for %v", artifact_path)
6450
continue
6551
}
66-
fd.Close()
6752

68-
reformatted, err := manager.ReformatVQL(ctx, string(data))
69-
if err != nil {
70-
returned_errs[artifact_path] = err
71-
continue
72-
}
53+
artifact_paths = append(artifact_paths, abs)
54+
}
7355

74-
out_fd, err := os.OpenFile(artifact_path,
75-
os.O_TRUNC|os.O_WRONLY|os.O_CREATE, 0600)
76-
if err != nil {
77-
returned_errs[artifact_path] = err
78-
continue
79-
}
80-
_, _ = out_fd.Write([]byte(reformatted))
81-
out_fd.Close()
56+
artifact_logger := &LogWriter{config_obj: sm.Config}
57+
builder := services.ScopeBuilder{
58+
Config: sm.Config,
59+
ACLManager: acl_managers.NewRoleACLManager(sm.Config, "administrator"),
60+
Logger: log.New(artifact_logger, "", 0),
61+
Env: ordereddict.NewDict().
62+
Set("DryRun", *reformat_dry_run).
63+
Set("Artifacts", artifact_paths),
8264
}
8365

84-
var ret error
85-
for artifact_path, err := range returned_errs {
86-
if err != nil {
87-
logger.Error("%v: <red>%v</>", artifact_path, err)
88-
ret = err
89-
} else {
90-
logger.Info("Reformatted %v: <green>OK</>", artifact_path)
91-
}
66+
query := `
67+
SELECT reformat(artifact=read_file(filename=_value)) AS Result
68+
FROM foreach(row=Artifacts)
69+
WHERE if(condition=Result.Error,
70+
then=log(level="ERROR", message="%v: <red>%v</>",
71+
args=[_value, Result.Error], dedup=-1),
72+
else=log(message="Reformatted %v: <green>OK</>",
73+
args=_value, dedup=-1)
74+
AND NOT DryRun
75+
AND copy(accessor="data", dest=_value, filename=Result.Artifact))
76+
AND FALSE
77+
`
78+
err = runQueryWithEnv(query, builder, "json")
79+
if err != nil {
80+
logger.Error("reformat: error running query: %v", query)
9281
}
9382

94-
return ret
83+
return artifact_logger.Error
9584
}
9685

9786
func init() {

‎docs/references/vql.yaml‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8621,6 +8621,23 @@
86218621
- linux_amd64_cgo
86228622
- windows_386_cgo
86238623
- windows_amd64_cgo
8624+
- name: reformat
8625+
description: |-
8626+
Reformat VQL
8627+
8628+
This function will reformat the artifact provided and return the reformatted content.
8629+
type: Function
8630+
args:
8631+
- name: artifact
8632+
type: string
8633+
description: The artifact VQL to reformat.
8634+
required: true
8635+
platforms:
8636+
- darwin_amd64_cgo
8637+
- darwin_arm64_cgo
8638+
- linux_amd64_cgo
8639+
- windows_386_cgo
8640+
- windows_amd64_cgo
86248641
- name: reg
86258642
description: |
86268643
An alias for the `registry` accessor, which accesses the registry using the

‎vql/functions/reformat.go‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
package functions
2+
3+
import (
4+
"context"
5+
"strings"
6+
7+
"github.com/Velocidex/ordereddict"
8+
"www.velocidex.com/golang/velociraptor/services"
9+
vql_subsystem "www.velocidex.com/golang/velociraptor/vql"
10+
"www.velocidex.com/golang/vfilter"
11+
"www.velocidex.com/golang/vfilter/arg_parser"
12+
)
13+
14+
type ReformatFunctionArgs struct {
15+
Artifact string `vfilter:"required,field=artifact,doc=The artifact VQL to reformat."`
16+
}
17+
18+
type ReformatFunction struct{}
19+
20+
type ReformatFunctionResult struct {
21+
Artifact string
22+
Error string
23+
}
24+
25+
func (self *ReformatFunctionResult) ToDict() *ordereddict.Dict {
26+
return ordereddict.NewDict().
27+
Set("Artifact", self.Artifact).
28+
Set("Error", self.Error)
29+
}
30+
31+
func (self *ReformatFunction) Call(ctx context.Context, scope vfilter.Scope, args *ordereddict.Dict) vfilter.Any {
32+
defer vql_subsystem.RegisterMonitor(ctx, "reformat", args)()
33+
34+
result := &ReformatFunctionResult{}
35+
36+
arg := &ReformatFunctionArgs{}
37+
err := arg_parser.ExtractArgsWithContext(ctx, scope, args, arg)
38+
if err != nil {
39+
result.Artifact = arg.Artifact
40+
result.Error = err.Error()
41+
return result.ToDict()
42+
}
43+
44+
config_obj, ok := vql_subsystem.GetServerConfig(scope)
45+
if !ok {
46+
scope.Log("reformat: Must be run on the server")
47+
return vfilter.Null{}
48+
}
49+
50+
manager, err := services.GetRepositoryManager(config_obj)
51+
if err != nil {
52+
result.Artifact = arg.Artifact
53+
result.Error = err.Error()
54+
return result.ToDict()
55+
}
56+
57+
reformatted, err := manager.ReformatVQL(ctx, arg.Artifact)
58+
if err != nil {
59+
result.Artifact = arg.Artifact
60+
result.Error = err.Error()
61+
return result.ToDict()
62+
}
63+
64+
result.Artifact = strings.Trim(reformatted, "\n")
65+
result.Error = ""
66+
67+
return result.ToDict()
68+
}
69+
70+
func (self ReformatFunction) Info(scope vfilter.Scope, type_map *vfilter.TypeMap) *vfilter.FunctionInfo {
71+
return &vfilter.FunctionInfo{
72+
Name: "reformat",
73+
Doc: `Reformat VQL
74+
75+
This function will reformat the artifact provided and return the reformatted content.`,
76+
ArgType: type_map.AddType(scope, &ReformatFunctionArgs{}),
77+
}
78+
}
79+
80+
func init() {
81+
vql_subsystem.RegisterFunction(&ReformatFunction{})
82+
}

0 commit comments

Comments
 (0)