Skip to content

Commit 313afdf

Browse files
joseotorojoseotoro
authored
Add JSON/CSV output support (#69)
1 parent 86b5069 commit 313afdf

File tree

8 files changed

+517
-20
lines changed

8 files changed

+517
-20
lines changed

‎README.md‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,11 @@ Restart the shell.
116116
$ vt file 8739c76e681f900923b900c9df0ef75cf421d39cabb54650c4b9ad19b6a76d85
117117
```
118118

119+
* Get information about a file in JSON format:
120+
```
121+
$ vt file 8739c76e681f900923b900c9df0ef75cf421d39cabb54650c4b9ad19b6a76d85 --format json
122+
```
123+
119124
* Get a specific analysis report for a file:
120125
```
121126
$ # File analysis IDs can be given as `f-<file_SHA256_hash>-<UNIX timestamp>`...
@@ -177,6 +182,16 @@ Restart the shell.
177182
status: "queued"
178183
```
179184

185+
* Export detections and tags of files from a search in CSV format:
186+
```
187+
$ vt search "positives:5+ type:pdf" -i sha256,last_analysis_stats.malicious,tags --format csv
188+
```
189+
190+
* Export detections and tags of files from a search in JSON format:
191+
```
192+
$ vt search "positives:5+ type:pdf" -i sha256,last_analysis_stats.malicious,tags --format json
193+
```
194+
180195
## Getting only what you want
181196

182197
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‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,12 @@ func addAPIKeyFlag(flags *pflag.FlagSet) {
3838
"API key")
3939
}
4040

41+
func addFormatFlag(flags *pflag.FlagSet) {
42+
flags.String(
43+
"format", "yaml",
44+
"Output format (yaml/json/csv)")
45+
}
46+
4147
func addHostFlag(flags *pflag.FlagSet) {
4248
flags.String(
4349
"host", "www.virustotal.com",

‎cmd/vt.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ func NewVTCommand() *cobra.Command {
5959
}
6060

6161
addAPIKeyFlag(cmd.PersistentFlags())
62+
addFormatFlag(cmd.PersistentFlags())
6263
addHostFlag(cmd.PersistentFlags())
6364
addProxyFlag(cmd.PersistentFlags())
6465
addVerboseFlag(cmd.PersistentFlags())

‎csv/csv.go‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
// Copyright © 2023 The VirusTotal CLI authors. All Rights Reserved.
2+
// Licensed under the Apache License, Version 2.0 (the "License");
3+
// you may not use this file except in compliance with the License.
4+
// You may obtain a copy of the License at
5+
//
6+
// http://www.apache.org/licenses/LICENSE-2.0
7+
//
8+
// Unless required by applicable law or agreed to in writing, software
9+
// distributed under the License is distributed on an "AS IS" BASIS,
10+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
// See the License for the specific language governing permissions and
12+
// limitations under the License.
13+
14+
package csv
15+
16+
import (
17+
"encoding/csv"
18+
"fmt"
19+
"io"
20+
"reflect"
21+
"sort"
22+
)
23+
24+
// An Encoder writes values as CSV to an output stream.
25+
type Encoder struct {
26+
w io.Writer
27+
}
28+
29+
// NewEncoder returns a new CSV encoder that writes to w.
30+
func NewEncoder(w io.Writer) *Encoder {
31+
return &Encoder{w: w}
32+
}
33+
34+
// Encode writes the CSV encoding of v to the stream.
35+
func (enc *Encoder) Encode(v interface{}) error {
36+
if v == nil {
37+
_, err := enc.w.Write([]byte("null"))
38+
return err
39+
}
40+
41+
var items []interface{}
42+
val := reflect.ValueOf(v)
43+
switch val.Kind() {
44+
case reflect.Slice:
45+
items = make([]interface{}, val.Len())
46+
for i := 0; i < val.Len(); i++ {
47+
items[i] = val.Index(i).Interface()
48+
}
49+
default:
50+
items = []interface{}{v}
51+
}
52+
numObjects := len(items)
53+
flattenObjects := make([]map[string]interface{}, numObjects)
54+
for i := 0; i < numObjects; i++ {
55+
f, err := flatten(items[i])
56+
if err != nil {
57+
return err
58+
}
59+
flattenObjects[i] = f
60+
}
61+
62+
keys := make(map[string]struct{})
63+
for _, o := range flattenObjects {
64+
for k := range o {
65+
keys[k] = struct{}{}
66+
}
67+
}
68+
69+
header := make([]string, len(keys))
70+
i := 0
71+
for k := range keys {
72+
header[i] = k
73+
i++
74+
}
75+
sort.Strings(header)
76+
77+
w := csv.NewWriter(enc.w)
78+
if len(header) > 1 || len(header) == 0 && header[0] != "" {
79+
if err := w.Write(header); err != nil {
80+
return err
81+
}
82+
}
83+
84+
for _, o := range flattenObjects {
85+
record := make([]string, len(keys))
86+
for i, key := range header {
87+
val, ok := o[key]
88+
if ok && val != nil {
89+
record[i] = fmt.Sprintf("%v", val)
90+
}
91+
}
92+
if err := w.Write(record); err != nil {
93+
return err
94+
}
95+
}
96+
w.Flush()
97+
return w.Error()
98+
}

‎csv/csv_test.go‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
// Copyright © 2023 The VirusTotal CLI authors. All Rights Reserved.
2+
// Licensed under the Apache License, Version 2.0 (the "License");
3+
// you may not use this file except in compliance with the License.
4+
// You may obtain a copy of the License at
5+
//
6+
// http://www.apache.org/licenses/LICENSE-2.0
7+
//
8+
// Unless required by applicable law or agreed to in writing, software
9+
// distributed under the License is distributed on an "AS IS" BASIS,
10+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
// See the License for the specific language governing permissions and
12+
// limitations under the License.
13+
14+
package csv
15+
16+
import (
17+
"bytes"
18+
"testing"
19+
20+
"github.com/stretchr/testify/assert"
21+
)
22+
23+
type Case struct {
24+
data interface{}
25+
expected string
26+
}
27+
28+
var csvTests = []Case{
29+
{
30+
data: nil,
31+
expected: "null",
32+
},
33+
{
34+
data: []int{1, 2, 3},
35+
expected: "1\n2\n3\n",
36+
},
37+
{
38+
data: map[string]interface{}{
39+
"b": []int{1, 2},
40+
"a": 2,
41+
"c": nil,
42+
},
43+
expected: "a,b,c\n2,\"1,2\",null\n",
44+
},
45+
}
46+
47+
func TestCSV(t *testing.T) {
48+
for _, test := range csvTests {
49+
b := new(bytes.Buffer)
50+
err := NewEncoder(b).Encode(test.data)
51+
assert.NoError(t, err)
52+
assert.Equal(t, test.expected, b.String(), "Test %v", test.data)
53+
}
54+
}

‎csv/flatten.go‎

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
// Copyright © 2023 The VirusTotal CLI authors. All Rights Reserved.
2+
// Licensed under the Apache License, Version 2.0 (the "License");
3+
// you may not use this file except in compliance with the License.
4+
// You may obtain a copy of the License at
5+
//
6+
// http://www.apache.org/licenses/LICENSE-2.0
7+
//
8+
// Unless required by applicable law or agreed to in writing, software
9+
// distributed under the License is distributed on an "AS IS" BASIS,
10+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
// See the License for the specific language governing permissions and
12+
// limitations under the License.
13+
14+
package csv
15+
16+
import (
17+
"encoding/json"
18+
"fmt"
19+
"reflect"
20+
"strings"
21+
)
22+
23+
func flatten(i interface{}) (map[string]interface{}, error) {
24+
result := make(map[string]interface{})
25+
err := flattenValue(reflect.ValueOf(i), "", result)
26+
return result, err
27+
}
28+
29+
func flattenValue(v reflect.Value, prefix string, m map[string]interface{}) error {
30+
switch v.Kind() {
31+
case reflect.Map:
32+
return flattenMap(v, prefix, m)
33+
case reflect.Struct:
34+
return flattenStruct(v, prefix, m)
35+
case reflect.Slice:
36+
return flattenSlice(v, prefix, m)
37+
case reflect.Interface, reflect.Ptr:
38+
if v.IsNil() {
39+
m[prefix] = "null"
40+
} else {
41+
return flattenValue(v.Elem(), prefix, m)
42+
}
43+
default:
44+
m[prefix] = v.Interface()
45+
}
46+
return nil
47+
}
48+
49+
func flattenSlice(v reflect.Value, prefix string, m map[string]interface{}) error {
50+
n := v.Len()
51+
if n == 0 {
52+
return nil
53+
}
54+
55+
first := v.Index(0)
56+
if first.Kind() == reflect.Interface {
57+
if !first.IsNil() {
58+
first = first.Elem()
59+
}
60+
}
61+
62+
switch first.Kind() {
63+
case reflect.Map, reflect.Slice, reflect.Struct:
64+
// Add the JSON representation of lists with complex types.
65+
// Otherwise the number of CSV headers can grow significantly.
66+
b, err := json.Marshal(v.Interface())
67+
if err != nil {
68+
return err
69+
}
70+
m[prefix] = string(b)
71+
default:
72+
values := make([]string, v.Len())
73+
for i := 0; i < v.Len(); i++ {
74+
val := v.Index(i).Interface()
75+
if val == nil {
76+
values[i] = "null"
77+
} else {
78+
values[i] = fmt.Sprintf("%v", val)
79+
}
80+
}
81+
m[prefix] = strings.Join(values, ",")
82+
}
83+
return nil
84+
}
85+
86+
func flattenStruct(v reflect.Value, prefix string, m map[string]interface{}) (err error) {
87+
n := v.NumField()
88+
if prefix != "" {
89+
prefix += "/"
90+
}
91+
for i := 0; i < n; i++ {
92+
typeField := v.Type().Field(i)
93+
key := typeField.Tag.Get("csv")
94+
if key == "" {
95+
key = v.Type().Field(i).Name
96+
}
97+
if err = flattenValue(v.Field(i), prefix+key, m); err != nil {
98+
return err
99+
}
100+
}
101+
return err
102+
}
103+
104+
func flattenMap(v reflect.Value, prefix string, m map[string]interface{}) (err error) {
105+
if prefix != "" {
106+
prefix += "/"
107+
}
108+
for _, k := range v.MapKeys() {
109+
if err := flattenValue(v.MapIndex(k), fmt.Sprintf("%v%v", prefix, k.Interface()), m); err != nil {
110+
return err
111+
}
112+
}
113+
return nil
114+
}

0 commit comments

Comments
 (0)