Skip to content

Commit

Permalink
Support addition on dicts (#1785)
Browse files Browse the repository at this point in the history
* Support addition on dicts

This results in a merge of the two dicts

dict(A=1, B=2) + dict(A=3, X=1) -> {B:2,A:3,X:1}

* Hash VQL function should return a dict.
  • Loading branch information
scudette authored May 10, 2022
1 parent 4c65fea commit 4ec9267
Show file tree
Hide file tree
Showing 6 changed files with 120 additions and 45 deletions.
8 changes: 6 additions & 2 deletions artifacts/testdata/server/testcases/functions.in.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,10 @@ Queries:
})
}) AS Filtered FROM scope()

# Test dict addition
- SELECT dict(A=1, B=2) + dict(A=3, X=1)
FROM scope()

# Test cache functions (first 2 should be same value due to caching)
- LET Foo(X) = if(condition=log(message=format(format="I Ran with %v", args=X)),
then=X + 5)
Expand Down Expand Up @@ -62,6 +66,6 @@ Queries:
FROM scope()

# Test entropy function
- SELECT
- SELECT
entropy(string="Hello world")
FROM scope()
FROM scope()
8 changes: 8 additions & 0 deletions artifacts/testdata/server/testcases/functions.out.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,14 @@ LET rows <= SELECT * FROM scope()[]SELECT len(list=["a", "b"]), len(list="hello"
"Z": 3
}
}
]SELECT dict(A=1, B=2) + dict(A=3, X=1) FROM scope()[
{
"dict(A=1, B=2) + dict(A=3, X=1)": {
"B": 2,
"A": 3,
"X": 1
}
}
]LET Foo(X) = if(condition=log(message=format(format="I Ran with %v", args=X)), then=X + 5)[]SELECT cache(func=Foo(X=5), key=5), cache(func=Foo(X=10), key=5), cache(func=Foo(X=10), key=10) FROM scope()[
{
"cache(func=Foo(X=5), key=5)": 10,
Expand Down
14 changes: 12 additions & 2 deletions utils/nil.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,22 @@ import (
// We need to do this stupid check because Go does not allow
// comparison to nil with interfaces.
func IsNil(v interface{}) bool {
if v == nil {
return true
}

switch v.(type) {
case types.Null, *types.Null:
return true
}

switch reflect.TypeOf(v).Kind() {
case reflect.Ptr, reflect.Map, reflect.Chan, reflect.Slice:
//use of IsNil method
return reflect.ValueOf(v).IsNil()

default:
return v == nil || (reflect.ValueOf(v).Kind() == reflect.Ptr &&
reflect.ValueOf(v).IsNil())
return false
}
}

Expand Down
13 changes: 10 additions & 3 deletions vql/functions/hash.go
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,13 @@ type HashResult struct {
sha256 hash.Hash
}

func (self *HashResult) ToDict() *ordereddict.Dict {
return ordereddict.NewDict().
Set("MD5", self.MD5).
Set("SHA1", self.SHA1).
Set("SHA256", self.SHA256)
}

type HashFunctionArgs struct {
Path *accessors.OSPath `vfilter:"required,field=path,doc=Path to open and hash."`
Accessor string `vfilter:"optional,field=accessor,doc=The accessor to use"`
Expand Down Expand Up @@ -96,10 +103,10 @@ func (self *HashFunction) Call(ctx context.Context,
}
defer file.Close()

result := HashResult{}
result := &HashResult{}

if arg.HashSelect == nil {
result = HashResult{
result = &HashResult{
md5: md5.New(),
sha1: sha1.New(),
sha256: sha256.New(),
Expand Down Expand Up @@ -146,7 +153,7 @@ func (self *HashFunction) Call(ctx context.Context,
result.SHA256 = fmt.Sprintf(
"%x", result.sha256.Sum(nil))
}
return result
return result.ToDict()
}

} else if err != nil {
Expand Down
84 changes: 84 additions & 0 deletions vql/protocols/dict.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,84 @@
package protocols

import (
"context"
"reflect"

"github.com/Velocidex/ordereddict"
vql_subsystem "www.velocidex.com/golang/velociraptor/vql"
"www.velocidex.com/golang/vfilter"
"www.velocidex.com/golang/vfilter/types"
)

type _BoolDict struct{}

func (self _BoolDict) Applicable(a vfilter.Any) bool {
switch a.(type) {
case ordereddict.Dict, *ordereddict.Dict:
return true
}

rt := reflect.TypeOf(a)
if rt == nil {
return false
}
return rt.Kind() == reflect.Slice || rt.Kind() == reflect.Map
}

func (self _BoolDict) Bool(ctx context.Context, scope vfilter.Scope, a vfilter.Any) bool {
switch t := a.(type) {
case ordereddict.Dict:
return t.Len() > 0
case *ordereddict.Dict:
return t.Len() > 0
}

rt := reflect.TypeOf(a)
if rt == nil {
return false
}
if rt.Kind() == reflect.Slice || rt.Kind() == reflect.Map {
a_slice := reflect.ValueOf(a)
return a_slice.Len() > 0
}

return false
}

type _AddDict struct{}

func (self _AddDict) Applicable(a types.Any, b types.Any) bool {
_, a_ok := a.(*ordereddict.Dict)
_, b_ok := b.(*ordereddict.Dict)
return a_ok && b_ok
}

func (self _AddDict) Add(scope types.Scope, a types.Any, b types.Any) types.Any {
a_dict, a_ok := a.(*ordereddict.Dict)
if !a_ok {
return &vfilter.Null{}
}
b_dict, b_ok := b.(*ordereddict.Dict)
if !b_ok {
return &vfilter.Null{}
}

res := ordereddict.NewDict()

for _, k := range a_dict.Keys() {
v, _ := a_dict.Get(k)
res.Set(k, v)
}

for _, k := range b_dict.Keys() {
v, _ := b_dict.Get(k)
res.Set(k, v)
}

return res
}

func init() {
vql_subsystem.RegisterProtocol(&_BoolDict{})
vql_subsystem.RegisterProtocol(&_AddDict{})
}
38 changes: 0 additions & 38 deletions vql/protocols/protocols.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,52 +2,15 @@ package protocols

import (
"context"
"reflect"
"time"

"github.com/Velocidex/ordereddict"
"www.velocidex.com/golang/velociraptor/accessors"
vql_subsystem "www.velocidex.com/golang/velociraptor/vql"
"www.velocidex.com/golang/vfilter"
"www.velocidex.com/golang/vfilter/protocols"
"www.velocidex.com/golang/vfilter/types"
)

type _BoolDict struct{}

func (self _BoolDict) Applicable(a vfilter.Any) bool {
switch a.(type) {
case ordereddict.Dict, *ordereddict.Dict:
return true
}

rt := reflect.TypeOf(a)
if rt == nil {
return false
}
return rt.Kind() == reflect.Slice || rt.Kind() == reflect.Map
}

func (self _BoolDict) Bool(ctx context.Context, scope vfilter.Scope, a vfilter.Any) bool {
switch t := a.(type) {
case ordereddict.Dict:
return t.Len() > 0
case *ordereddict.Dict:
return t.Len() > 0
}

rt := reflect.TypeOf(a)
if rt == nil {
return false
}
if rt.Kind() == reflect.Slice || rt.Kind() == reflect.Map {
a_slice := reflect.ValueOf(a)
return a_slice.Len() > 0
}

return false
}

type _BoolTime struct{}

func (self _BoolTime) Applicable(a vfilter.Any) bool {
Expand Down Expand Up @@ -135,7 +98,6 @@ func (self _GlobFileInfoAssociative) GetMembers(
}

func init() {
vql_subsystem.RegisterProtocol(&_BoolDict{})
vql_subsystem.RegisterProtocol(&_BoolTime{})
vql_subsystem.RegisterProtocol(&_BoolEq{})
vql_subsystem.RegisterProtocol(&_GlobFileInfoAssociative{})
Expand Down

0 comments on commit 4ec9267

Please sign in to comment.