Add JSON/CSV output support (#69)

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-<file_SHA256_hash>-<UNIX timestamp>`...
@@ -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).

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",

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())

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()
+}

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)
+	}
+}

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
+}