Improve filtering with --include and --exclude.

Closes #31.

Author: plusvic

cmd/relationship.go

@@ -17,11 +17,11 @@ import (
 	"encoding/base64"
 	"encoding/gob"
 	"fmt"
+	"github.com/VirusTotal/vt-cli/utils"
 	"os"
 	"path"
 	"sync"
 
-	"github.com/VirusTotal/vt-cli/utils"
 	vt "github.com/VirusTotal/vt-go"
 	homedir "github.com/mitchellh/go-homedir"
 	"github.com/spf13/cobra"
@@ -106,7 +106,7 @@ func NewRelationshipsCmd(collection, objectType, use string) *cobra.Command {
 		Args:  cobra.ExactArgs(1),
 		RunE: func(cmd *cobra.Command, args []string) error {
 			var wg sync.WaitGroup
-			var m sync.Map
+			var sm sync.Map
 			for _, r := range objectRelationshipsMap[objectType] {
 				wg.Add(1)
 				go func(relationshipName string) {
@@ -118,17 +118,30 @@ func NewRelationshipsCmd(collection, objectType, use string) *cobra.Command {
 					if err != nil {
 						fmt.Println(err)
 					} else if len(objs) > 0 {
-						m.Store(relationshipName, objs)
+						sm.Store(relationshipName, objs)
 					}
 					wg.Done()
 				}(r.Name)
 			}
 			wg.Wait()
+
+			m := make(map[string]interface{})
+			sm.Range(func(key, value interface{}) bool {
+				m[key.(string)] = value
+				return true
+			})
+
+			if viper.IsSet("include") || viper.IsSet("exclude") {
+				m = utils.FilterMap(m,
+					viper.GetStringSlice("include"),
+					viper.GetStringSlice("exclude"))
+			}
+
 			p, err := NewPrinter(cmd)
 			if err != nil {
 				return err
 			}
-			return p.PrintSyncMap(&m)
+			return p.Print(m)
 		},
 	}
 

utils/filter.go

@@ -14,14 +14,14 @@
 package utils
 
 import (
-	"reflect"
-
 	glob "github.com/gobwas/glob"
+	"reflect"
+	"strings"
 )
 
-// FilterMap receives a map with string keys and arbirary values (possibly
+// FilterMap receives a map with string keys and arbitrary values (possibly
 // other maps) and return a new map which is a subset of the original one
-// contaning only the keys matching any of the patterns in "include" and
+// containing only the keys matching any of the patterns in "include" and
 // excluding keys matching any of the patterns in "exclude". The logic for
 // determining if a key matches the pattern goes as follow:
 //
@@ -45,10 +45,27 @@ func FilterMap(m map[string]interface{}, include, exclude []string) map[string]i
 		cp := glob.MustCompile(p, '.')
 		includeGlob[i] = cp
 	}
+	// For each include pattern that do not ends with **, add the same pattern
+	// but ended in .**. This because when someone says that she wants to include
+	// "foo.bar", where "foo.bar" is dictionary, what she actually expects is
+	// getting the dictionary with all its keys, but the keys inside the
+	// dictionary don't match the "foo.bar" pattern, so we add "foo.bar.**".
+	for _, p := range include {
+		if !strings.HasSuffix(p, "**") {
+			includeGlob = append(includeGlob, glob.MustCompile(p+".**", '.'))
+		}
+	}
 	for i, p := range exclude {
 		cp := glob.MustCompile(p, '.')
 		excludeGlob[i] = cp
 	}
+	// The same happens if you exclude "foo.bar", what you actually mean is
+	// excluding "foo.bar" and "foo.bar.**".
+	for _, p := range exclude {
+		if !strings.HasSuffix(p, "**") {
+			excludeGlob = append(excludeGlob, glob.MustCompile(p+".**", '.'))
+		}
+	}
 	filtered := filterMap(reflect.ValueOf(m), includeGlob, excludeGlob, "")
 	return filtered.Interface().(map[string]interface{})
 }

utils/filter_test.go

@@ -73,6 +73,31 @@ var testCases = []testCase{
 		},
 	},
 
+	testCase{
+		include: []string{"foo"},
+		exclude: []string{"**.quux"},
+		input:   testMap,
+		output: map[string]interface{}{
+
+		},
+	},
+
+	testCase{
+		include: []string{"foo**"},
+		exclude: []string{"**.key1"},
+		input:   testMap,
+		output: map[string]interface{}{
+			"foo": map[string]interface{}{
+				"qux": map[string]interface{}{
+					"quux": map[string]interface{}{
+						"key2": "val2",
+						"key3": []string{"val3"},
+					},
+				},
+			},
+		},
+	},
+
 	testCase{
 		include: []string{"foo**"},
 		exclude: []string{"**.key1"},
@@ -179,6 +204,21 @@ var testCases = []testCase{
 		},
 	},
 
+	testCase{
+		include: []string{"qux"},
+		input:   testMap,
+		output: map[string]interface{}{
+			"qux": []interface{}{
+				map[string]interface{}{
+					"key4": "val4",
+				},
+				map[string]interface{}{
+					"key5": "val5",
+				},
+			},
+		},
+	},
+
 	testCase{
 		include: []string{"ary"},
 		input:   testMap,

utils/printer.go

@@ -63,6 +63,11 @@ func (p *Printer) PrintSyncMap(sm *sync.Map) error {
 		m[key.(string)] = value
 		return true
 	})
+	if viper.IsSet("include") || viper.IsSet("exclude") {
+		m = FilterMap(m,
+			viper.GetStringSlice("include"),
+			viper.GetStringSlice("exclude"))
+	}
 	return p.Print(m)
 }
 
@@ -95,21 +100,27 @@ func ObjectToMap(obj *vt.Object) map[string]interface{} {
 			m[name] = l
 		}
 	}
-	if viper.IsSet("include") || viper.IsSet("exclude") {
-		m = FilterMap(m,
-			viper.GetStringSlice("include"),
-			viper.GetStringSlice("exclude"))
-	}
 	return m
 }
 
 // PrintObjects prints all the specified objects to stdout.
 func (p *Printer) PrintObjects(objs []*vt.Object) error {
 	list := make([]map[string]interface{}, 0)
 	for _, obj := range objs {
-		list = append(list, ObjectToMap(obj))
+		m := ObjectToMap(obj)
+		if viper.IsSet("include") || viper.IsSet("exclude") {
+			m = FilterMap(m,
+				viper.GetStringSlice("include"),
+				viper.GetStringSlice("exclude"))
+		}
+		if len(m) > 0 {
+			list = append(list, m)
+		}
+	}
+	if len(list) > 0 {
+		return p.Print(list)
 	}
-	return p.Print(list)
+	return nil
 }
 
 // PrintObject prints the specified object to stdout.