diff --git a/README.md b/README.md index d1aa27b..b94dab5 100644 --- a/README.md +++ b/README.md @@ -116,6 +116,11 @@ Restart the shell. $ vt file 8739c76e681f900923b900c9df0ef75cf421d39cabb54650c4b9ad19b6a76d85 ``` +* Get information about a file in JSON format: + ``` + $ vt file 8739c76e681f900923b900c9df0ef75cf421d39cabb54650c4b9ad19b6a76d85 --format json + ``` + * Get a specific analysis report for a file: ``` $ # File analysis IDs can be given as `f--`... @@ -177,6 +182,16 @@ Restart the shell. status: "queued" ``` +* Export detections and tags of files from a search in CSV format: + ``` + $ vt search "positives:5+ type:pdf" -i sha256,last_analysis_stats.malicious,tags --format csv + ``` + +* Export detections and tags of files from a search in JSON format: + ``` + $ vt search "positives:5+ type:pdf" -i sha256,last_analysis_stats.malicious,tags --format json + ``` + ## Getting only what you want When you ask for information about a file, URL, domain, IP address or any other object in VirusTotal, you get a lot of data (by default in YAML format) that is usually more than what you need. You can narrow down the information shown by the vt-cli tool by using the `--include` and `--exclude` command-line options (`-i` and `-x` in short form). diff --git a/cmd/cmd.go b/cmd/cmd.go index cd7a0da..61f5e55 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -38,6 +38,12 @@ func addAPIKeyFlag(flags *pflag.FlagSet) { "API key") } +func addFormatFlag(flags *pflag.FlagSet) { + flags.String( + "format", "yaml", + "Output format (yaml/json/csv)") +} + func addHostFlag(flags *pflag.FlagSet) { flags.String( "host", "www.virustotal.com", diff --git a/cmd/vt.go b/cmd/vt.go index d809858..61e41fd 100644 --- a/cmd/vt.go +++ b/cmd/vt.go @@ -59,6 +59,7 @@ func NewVTCommand() *cobra.Command { } addAPIKeyFlag(cmd.PersistentFlags()) + addFormatFlag(cmd.PersistentFlags()) addHostFlag(cmd.PersistentFlags()) addProxyFlag(cmd.PersistentFlags()) addVerboseFlag(cmd.PersistentFlags()) diff --git a/csv/csv.go b/csv/csv.go new file mode 100644 index 0000000..b93ae16 --- /dev/null +++ b/csv/csv.go @@ -0,0 +1,98 @@ +// Copyright © 2023 The VirusTotal CLI authors. All Rights Reserved. +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package csv + +import ( + "encoding/csv" + "fmt" + "io" + "reflect" + "sort" +) + +// An Encoder writes values as CSV to an output stream. +type Encoder struct { + w io.Writer +} + +// NewEncoder returns a new CSV encoder that writes to w. +func NewEncoder(w io.Writer) *Encoder { + return &Encoder{w: w} +} + +// Encode writes the CSV encoding of v to the stream. +func (enc *Encoder) Encode(v interface{}) error { + if v == nil { + _, err := enc.w.Write([]byte("null")) + return err + } + + var items []interface{} + val := reflect.ValueOf(v) + switch val.Kind() { + case reflect.Slice: + items = make([]interface{}, val.Len()) + for i := 0; i < val.Len(); i++ { + items[i] = val.Index(i).Interface() + } + default: + items = []interface{}{v} + } + numObjects := len(items) + flattenObjects := make([]map[string]interface{}, numObjects) + for i := 0; i < numObjects; i++ { + f, err := flatten(items[i]) + if err != nil { + return err + } + flattenObjects[i] = f + } + + keys := make(map[string]struct{}) + for _, o := range flattenObjects { + for k := range o { + keys[k] = struct{}{} + } + } + + header := make([]string, len(keys)) + i := 0 + for k := range keys { + header[i] = k + i++ + } + sort.Strings(header) + + w := csv.NewWriter(enc.w) + if len(header) > 1 || len(header) == 0 && header[0] != "" { + if err := w.Write(header); err != nil { + return err + } + } + + for _, o := range flattenObjects { + record := make([]string, len(keys)) + for i, key := range header { + val, ok := o[key] + if ok && val != nil { + record[i] = fmt.Sprintf("%v", val) + } + } + if err := w.Write(record); err != nil { + return err + } + } + w.Flush() + return w.Error() +} diff --git a/csv/csv_test.go b/csv/csv_test.go new file mode 100644 index 0000000..77202bc --- /dev/null +++ b/csv/csv_test.go @@ -0,0 +1,54 @@ +// Copyright © 2023 The VirusTotal CLI authors. All Rights Reserved. +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package csv + +import ( + "bytes" + "testing" + + "github.com/stretchr/testify/assert" +) + +type Case struct { + data interface{} + expected string +} + +var csvTests = []Case{ + { + data: nil, + expected: "null", + }, + { + data: []int{1, 2, 3}, + expected: "1\n2\n3\n", + }, + { + data: map[string]interface{}{ + "b": []int{1, 2}, + "a": 2, + "c": nil, + }, + expected: "a,b,c\n2,\"1,2\",null\n", + }, +} + +func TestCSV(t *testing.T) { + for _, test := range csvTests { + b := new(bytes.Buffer) + err := NewEncoder(b).Encode(test.data) + assert.NoError(t, err) + assert.Equal(t, test.expected, b.String(), "Test %v", test.data) + } +} diff --git a/csv/flatten.go b/csv/flatten.go new file mode 100644 index 0000000..3f27e3c --- /dev/null +++ b/csv/flatten.go @@ -0,0 +1,114 @@ +// Copyright © 2023 The VirusTotal CLI authors. All Rights Reserved. +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package csv + +import ( + "encoding/json" + "fmt" + "reflect" + "strings" +) + +func flatten(i interface{}) (map[string]interface{}, error) { + result := make(map[string]interface{}) + err := flattenValue(reflect.ValueOf(i), "", result) + return result, err +} + +func flattenValue(v reflect.Value, prefix string, m map[string]interface{}) error { + switch v.Kind() { + case reflect.Map: + return flattenMap(v, prefix, m) + case reflect.Struct: + return flattenStruct(v, prefix, m) + case reflect.Slice: + return flattenSlice(v, prefix, m) + case reflect.Interface, reflect.Ptr: + if v.IsNil() { + m[prefix] = "null" + } else { + return flattenValue(v.Elem(), prefix, m) + } + default: + m[prefix] = v.Interface() + } + return nil +} + +func flattenSlice(v reflect.Value, prefix string, m map[string]interface{}) error { + n := v.Len() + if n == 0 { + return nil + } + + first := v.Index(0) + if first.Kind() == reflect.Interface { + if !first.IsNil() { + first = first.Elem() + } + } + + switch first.Kind() { + case reflect.Map, reflect.Slice, reflect.Struct: + // Add the JSON representation of lists with complex types. + // Otherwise the number of CSV headers can grow significantly. + b, err := json.Marshal(v.Interface()) + if err != nil { + return err + } + m[prefix] = string(b) + default: + values := make([]string, v.Len()) + for i := 0; i < v.Len(); i++ { + val := v.Index(i).Interface() + if val == nil { + values[i] = "null" + } else { + values[i] = fmt.Sprintf("%v", val) + } + } + m[prefix] = strings.Join(values, ",") + } + return nil +} + +func flattenStruct(v reflect.Value, prefix string, m map[string]interface{}) (err error) { + n := v.NumField() + if prefix != "" { + prefix += "/" + } + for i := 0; i < n; i++ { + typeField := v.Type().Field(i) + key := typeField.Tag.Get("csv") + if key == "" { + key = v.Type().Field(i).Name + } + if err = flattenValue(v.Field(i), prefix+key, m); err != nil { + return err + } + } + return err +} + +func flattenMap(v reflect.Value, prefix string, m map[string]interface{}) (err error) { + if prefix != "" { + prefix += "/" + } + for _, k := range v.MapKeys() { + if err := flattenValue(v.MapIndex(k), fmt.Sprintf("%v%v", prefix, k.Interface()), m); err != nil { + return err + } + } + return nil +} diff --git a/csv/flatten_test.go b/csv/flatten_test.go new file mode 100644 index 0000000..d798fa9 --- /dev/null +++ b/csv/flatten_test.go @@ -0,0 +1,177 @@ +// Copyright © 2023 The VirusTotal CLI authors. All Rights Reserved. +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package csv + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +type T struct { + data interface{} + expected map[string]interface{} +} + +var tests = []T{ + { + data: "foo", + expected: map[string]interface{}{"": "foo"}, + }, + { + data: 1, + expected: map[string]interface{}{"": 1}, + }, + { + data: false, + expected: map[string]interface{}{"": false}, + }, + { + data: true, + expected: map[string]interface{}{"": true}, + }, + { + data: map[string]string{}, + expected: map[string]interface{}{}, + }, + { + data: map[string]map[string]string{ + "foo": {}, + }, + expected: map[string]interface{}{}, + }, + { + data: []string{}, + expected: map[string]interface{}{}, + }, + { + data: map[string]string{ + "uno": "1", + "dos": "2", + "tres": "3", + "": "", + "#foo": "foo", + "|foo": "foo", + "_foo": "foo", + }, + expected: map[string]interface{}{ + "": "", + "uno": "1", + "dos": "2", + "tres": "3", + "#foo": "foo", + "|foo": "foo", + "_foo": "foo", + }, + }, + { + data: []string{ + "uno", + "dos", + "tres", + }, + expected: map[string]interface{}{"": "uno,dos,tres"}, + }, + { + data: struct { + Foo string + Bar string + }{ + "uno", + "dos", + }, + expected: map[string]interface{}{ + "Foo": "uno", + "Bar": "dos", + }, + }, + { + data: struct { + Foo string + }{ + "uno\ndos", + }, + expected: map[string]interface{}{"Foo": "uno\ndos"}, + }, + { + data: map[string]interface{}{ + "numbers": []interface{}{ + map[string]string{ + "number": "1", + "numeral": "first", + }, + map[string]string{ + "number": "2", + "numeral": "second", + }, + }, + }, + expected: map[string]interface{}{ + "numbers": "[{\"number\":\"1\",\"numeral\":\"first\"},{\"number\":\"2\",\"numeral\":\"second\"}]", + }, + }, + { + data: struct { + A map[string]string + B map[string][]int + }{ + A: map[string]string{"1": "xx", "2": "yy"}, + B: map[string][]int{"hello": {1, 2}}, + }, + expected: map[string]interface{}{ + "A/1": "xx", + "A/2": "yy", + "B/hello": "1,2", + }, + }, + { + data: map[string]interface{}{ + "key1": struct { + A int + B bool `csv:"field"` + C []bool + }{ + A: 2, + B: true, + C: []bool{true, false}, + }, + "key2": map[interface{}]interface{}{ + 1: []string{"hello", "world"}, + 2: []string{}, + true: "test", + 2.1: map[string]string{ + "x": "x", + "y": "", + }, + }, + }, + expected: map[string]interface{}{ + "key1/A": 2, + "key1/field": true, + "key1/C": "true,false", + "key2/1": "hello,world", + "key2/true": "test", + "key2/2.1/x": "x", + "key2/2.1/y": "", + }, + }, +} + +func TestFlatten(t *testing.T) { + for _, test := range tests { + result, err := flatten(test.data) + assert.NoError(t, err) + assert.Equal(t, test.expected, result, "Test %v", test.data) + } +} diff --git a/utils/printer.go b/utils/printer.go index 2e5f759..13015ae 100644 --- a/utils/printer.go +++ b/utils/printer.go @@ -14,6 +14,8 @@ package utils import ( + "encoding/json" + "errors" "fmt" "net/url" "os" @@ -21,6 +23,7 @@ import ( "strings" "sync" + "github.com/VirusTotal/vt-cli/csv" "github.com/VirusTotal/vt-cli/yaml" vt "github.com/VirusTotal/vt-go" "github.com/fatih/color" @@ -45,15 +48,26 @@ func NewPrinter(client *APIClient, cmd *cobra.Command, colors *yaml.Colors) (*Pr // Print prints the provided data to stdout. func (p *Printer) Print(data interface{}) error { - return yaml.NewEncoder( - ansi.NewAnsiStdout(), - yaml.EncoderColors(p.colors), - yaml.EncoderDateKeys([]glob.Glob{ - glob.MustCompile("last_login"), - glob.MustCompile("user_since"), - glob.MustCompile("date"), - glob.MustCompile("*_date"), - })).Encode(data) + format := strings.ToLower(viper.GetString("format")) + if format == "" || format == "yaml" { + return yaml.NewEncoder( + ansi.NewAnsiStdout(), + yaml.EncoderColors(p.colors), + yaml.EncoderDateKeys([]glob.Glob{ + glob.MustCompile("last_login"), + glob.MustCompile("user_since"), + glob.MustCompile("date"), + glob.MustCompile("*_date"), + })).Encode(data) + } else if format == "json" { + encoder := json.NewEncoder(ansi.NewAnsiStdout()) + encoder.SetIndent("", " ") + return encoder.Encode(data) + } else if format == "csv" { + return csv.NewEncoder(ansi.NewAnsiStdout()).Encode(data) + } else { + return errors.New("unknown format") + } } // PrintSyncMap prints a sync.Map. @@ -159,13 +173,21 @@ func (p *Printer) GetAndPrintObjects(endpoint string, r StringReader, argRe *reg go p.client.RetrieveObjects(endpoint, filteredArgs, objectsCh, errorsCh) - for obj := range objectsCh { - if viper.GetBool("identifiers-only") { - fmt.Printf("%s\n", obj.ID()) - } else { - if err := p.PrintObject(obj); err != nil { - return err - } + if viper.GetBool("identifiers-only") { + var objectIds []string + for obj := range objectsCh { + objectIds = append(objectIds, obj.ID()) + } + if err := p.Print(objectIds); err != nil { + return err + } + } else { + var objects []*vt.Object + for obj := range objectsCh { + objects = append(objects, obj) + } + if err := p.PrintObjects(objects); err != nil { + return err } } @@ -191,19 +213,29 @@ func (p *Printer) PrintCollection(collection *url.URL) error { // PrintIterator prints the objects returned by an object iterator. func (p *Printer) PrintIterator(it *vt.Iterator) error { + var objs []*vt.Object + var ids []string for it.Next() { obj := it.Get() if viper.GetBool("identifiers-only") { - fmt.Printf("%s\n", obj.ID()) + ids = append(ids, obj.ID()) } else { - if err := p.PrintObject(obj); err != nil { - return err - } + objs = append(objs, obj) } } if err := it.Error(); err != nil { return err } + + if viper.GetBool("identifiers-only") { + if err := p.Print(ids); err != nil { + return err + } + } else { + if err := p.PrintObjects(objs); err != nil { + return err + } + } p.PrintCommandLineWithCursor(it) return nil }