/
vql_test.go
111 lines (97 loc) · 2.78 KB
/
vql_test.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
package actions_test
import (
"testing"
"github.com/alecthomas/assert"
"github.com/stretchr/testify/suite"
"www.velocidex.com/golang/velociraptor/actions"
actions_proto "www.velocidex.com/golang/velociraptor/actions/proto"
artifacts_proto "www.velocidex.com/golang/velociraptor/artifacts/proto"
"www.velocidex.com/golang/velociraptor/file_store/test_utils"
"www.velocidex.com/golang/velociraptor/responder"
)
type ClientVQLTestSuite struct {
test_utils.TestSuite
}
func (self *ClientVQLTestSuite) TestCPUThrottler() {
request := &actions_proto.VQLCollectorArgs{
Query: []*actions_proto.VQLRequest{
{
Name: "Query",
VQL: "SELECT 'Boo' FROM scope()",
},
},
}
// Query is not limited
resp := responder.TestResponder(self.ConfigObj)
actions.VQLClientAction{}.StartQuery(self.ConfigObj, self.Sm.Ctx, resp, request)
resp.Close()
assert.NotContains(self.T(), getLogs(resp), "Will throttle query")
// Query will now be limited
resp = responder.TestResponder(self.ConfigObj)
request.CpuLimit = 20
actions.VQLClientAction{}.StartQuery(self.ConfigObj, self.Sm.Ctx, resp, request)
resp.Close()
assert.Contains(self.T(), getLogs(resp), "Will throttle query")
}
// Make sure that dependent artifacts are properly used
func (self *ClientVQLTestSuite) TestDependentArtifacts() {
resp := responder.TestResponder(self.ConfigObj)
actions.VQLClientAction{}.StartQuery(self.ConfigObj, self.Sm.Ctx, resp,
&actions_proto.VQLCollectorArgs{
Query: []*actions_proto.VQLRequest{
{
Name: "Query",
VQL: "SELECT * FROM Artifact.Custom.Foo.Bar.Baz.A()",
},
},
Artifacts: []*artifacts_proto.Artifact{
{
Name: "Custom.Foo.Bar.Baz.A",
Sources: []*artifacts_proto.ArtifactSource{
{
Query: "SELECT * FROM Artifact.Custom.Foo.Bar.Baz.B()",
},
},
},
{
Name: "Custom.Foo.Bar.Baz.B",
Sources: []*artifacts_proto.ArtifactSource{
{
Query: "SELECT * FROM Artifact.Custom.Foo.Bar.Baz.C()",
},
},
},
{
Name: "Custom.Foo.Bar.Baz.C",
Sources: []*artifacts_proto.ArtifactSource{
{
Query: "SELECT 1 AS X FROM scope()",
},
},
},
},
})
assert.Equal(self.T(), "{\"X\":1,\"_Source\":\"Custom.Foo.Bar.Baz.A\"}\n", getVQLResponse(resp))
}
func getLogs(resp *responder.Responder) string {
result := ""
responses := responder.GetTestResponses(resp)
for _, item := range responses {
if item.LogMessage != nil {
result += item.LogMessage.Jsonl + "\n"
}
}
return result
}
func getVQLResponse(resp *responder.Responder) string {
responses := responder.GetTestResponses(resp)
for _, item := range responses {
if item.VQLResponse != nil {
return item.VQLResponse.JSONLResponse
}
}
return ""
}
func TestClientVQL(t *testing.T) {
suite.Run(t, &ClientVQLTestSuite{})
}