Skip to content

Commit 6707ddd

Browse files
joseotorojoseotoro
authored
Add support for collections (#39)
1 parent 88f1eae commit 6707ddd

File tree

5 files changed

+215
-3
lines changed

5 files changed

+215
-3
lines changed

‎cmd/collection.go

Lines changed: 208 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,208 @@
1+
package cmd
2+
3+
import (
4+
"fmt"
5+
"io"
6+
"strings"
7+
8+
"github.com/VirusTotal/vt-go"
9+
10+
"github.com/VirusTotal/vt-cli/utils"
11+
"github.com/spf13/cobra"
12+
"github.com/spf13/viper"
13+
)
14+
15+
var collectionCmdHelp = `Get information about one or more collections.
16+
17+
This command receives one or more collection IDs and returns information about
18+
them. The information for each collection is returned in the same order as the
19+
collections are passed to the command.
20+
21+
If the command receives a single hypen (-) the collection will be read from
22+
the standard input, one per line.`
23+
24+
var collectionCmdExample = ` vt collection malpedia_win_emotet
25+
vt collection malpedia_win_emotet alienvault_603eb1abdd4812819c64e197
26+
cat list_of_collections | vt collection -`
27+
28+
// NewCollectionCmd returns a new instance of the 'collection' command.
29+
func NewCollectionCmd() *cobra.Command {
30+
cmd := &cobra.Command{
31+
Use: "collection [collection]...",
32+
Short: "Get information about collections",
33+
Long: collectionCmdHelp,
34+
Example: collectionCmdExample,
35+
Args: cobra.MinimumNArgs(1),
36+
37+
RunE: func(cmd *cobra.Command, args []string) error {
38+
p, err := NewPrinter(cmd)
39+
if err != nil {
40+
return err
41+
}
42+
return p.GetAndPrintObjects(
43+
"collections/%s",
44+
utils.StringReaderFromCmdArgs(args),
45+
nil)
46+
},
47+
}
48+
49+
cmd.AddCommand(NewCollectionCreateCmd())
50+
cmd.AddCommand(NewCollectionRenameCmd())
51+
cmd.AddCommand(NewCollectionUpdateCmd())
52+
53+
addRelationshipCmds(cmd, "collections", "collection", "[collection]")
54+
addThreadsFlag(cmd.Flags())
55+
addIncludeExcludeFlags(cmd.Flags())
56+
addIDOnlyFlag(cmd.Flags())
57+
58+
return cmd
59+
}
60+
61+
var createCollectionCmdHelp = `Creates a collection from a list of IOCs.
62+
63+
This command receives one of more IoCs (sha256 hashes, URLs, domains, IP addresses)
64+
and creates a collection from them.
65+
66+
If the command receives a single hypen (-) the IoCs will be read from the
67+
standard input.`
68+
69+
var createCollectionExample = ` vt collection create -n [collection_name] www.example.com
70+
vt collection create -n [collection_name] www.example.com 8.8.8.8
71+
cat list_of_iocs | vt collection create -n [collection_name] -`
72+
73+
// NewCollectionCreateCmd returns a command for creating a collection.
74+
func NewCollectionCreateCmd() *cobra.Command {
75+
cmd := &cobra.Command{
76+
Use: "create [ioc]...",
77+
Short: "Create a collection.",
78+
Long: createCollectionCmdHelp,
79+
Example: createCollectionExample,
80+
Args: cobra.MinimumNArgs(1),
81+
82+
RunE: func(cmd *cobra.Command, args []string) error {
83+
c, err := NewAPIClient()
84+
if err != nil {
85+
return err
86+
}
87+
p, err := NewPrinter(cmd)
88+
if err != nil {
89+
return err
90+
}
91+
reader := utils.StringReaderFromCmdArgs(args)
92+
93+
collection := vt.NewObject("collection")
94+
collection.SetString("name", viper.GetString("name"))
95+
collection.SetData("raw_items", rawFromReader(reader))
96+
97+
if err := c.PostObject(vt.URL("collections"), collection); err != nil {
98+
return err
99+
}
100+
101+
if viper.GetBool("identifiers-only") {
102+
fmt.Printf("%s\n", collection.ID())
103+
} else {
104+
if err := p.PrintObject(collection); err != nil {
105+
return err
106+
}
107+
}
108+
109+
return nil
110+
},
111+
}
112+
113+
cmd.Flags().StringP(
114+
"name", "n", "",
115+
"Collection's name (required)")
116+
_ = cmd.MarkFlagRequired("name")
117+
addIncludeExcludeFlags(cmd.Flags())
118+
addIDOnlyFlag(cmd.Flags())
119+
120+
return cmd
121+
}
122+
123+
func patchCollection(id, attr string, value interface{}) error {
124+
client, err := NewAPIClient()
125+
if err != nil {
126+
return err
127+
}
128+
obj := vt.NewObjectWithID("collection", id)
129+
obj.Set(attr, value)
130+
return client.PatchObject(vt.URL("collections/%s", id), obj)
131+
}
132+
133+
// NewCollectionRenameCmd returns a command for renaming a collection.
134+
func NewCollectionRenameCmd() *cobra.Command {
135+
return &cobra.Command{
136+
Use: "rename [collection id] [name]",
137+
Short: "Rename collection.",
138+
Args: cobra.ExactArgs(2),
139+
140+
RunE: func(cmd *cobra.Command, args []string) error {
141+
return patchCollection(args[0], "name", args[1])
142+
},
143+
}
144+
}
145+
146+
var updateCollectionCmdHelp = `Adds new items to a collection.
147+
148+
This command receives a collection ID and one of more IoCs
149+
(sha256 hashes, URLs, domains, IP addresses) and adds them to the collection.
150+
151+
If the command receives a single hypen (-) the IoCs will be read from the
152+
standard input.`
153+
154+
var updateCollectionExample = ` vt collection update [collection id] www.example.com
155+
vt collection update [collection id] www.example.com 8.8.8.8
156+
cat list_of_iocs | vt collection update [collection id] -`
157+
158+
// NewCollectionUpdateCmd returns a command for adding new items to a collection.
159+
func NewCollectionUpdateCmd() *cobra.Command {
160+
return &cobra.Command{
161+
Use: "update [collection id] [ioc]...",
162+
Short: "Add new items to a collection.",
163+
Args: cobra.MinimumNArgs(2),
164+
Long: updateCollectionCmdHelp,
165+
Example: updateCollectionExample,
166+
167+
RunE: func(cmd *cobra.Command, args []string) error {
168+
c, err := NewAPIClient()
169+
if err != nil {
170+
return err
171+
}
172+
p, err := NewPrinter(cmd)
173+
if err != nil {
174+
return err
175+
}
176+
177+
collection := vt.NewObjectWithID("collection", args[0])
178+
reader := utils.StringReaderFromCmdArgs(args[1:])
179+
collection.SetData("raw_items", rawFromReader(reader))
180+
181+
if err := c.PatchObject(vt.URL("collections/%s", args[0]), collection); err != nil {
182+
return err
183+
}
184+
185+
if viper.GetBool("identifiers-only") {
186+
fmt.Printf("%s\n", collection.ID())
187+
} else {
188+
if err := p.PrintObject(collection); err != nil {
189+
return err
190+
}
191+
}
192+
193+
return nil
194+
},
195+
}
196+
}
197+
198+
func rawFromReader(reader utils.StringReader) string {
199+
var lines []string
200+
for {
201+
next, err := reader.ReadString()
202+
if err == io.EOF {
203+
break
204+
}
205+
lines = append(lines, next)
206+
}
207+
return strings.Join(lines, " ")
208+
}

‎cmd/vt.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ func NewVTCommand() *cobra.Command {
6464
addVerboseFlag(cmd.PersistentFlags())
6565

6666
cmd.AddCommand(NewAnalysisCmd())
67+
cmd.AddCommand(NewCollectionCmd())
6768
cmd.AddCommand(NewCompletionCmd())
6869
cmd.AddCommand(NewDomainCmd())
6970
cmd.AddCommand(NewDownloadCmd())

‎go.mod

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ module github.com/VirusTotal/vt-cli
33
go 1.14
44

55
require (
6-
github.com/VirusTotal/vt-go v0.0.0-20211116094520-07a92e6467b7
6+
github.com/VirusTotal/vt-go v0.0.0-20211209151516-855a1e790678
77
github.com/briandowns/spinner v1.7.0
88
github.com/cavaliercoder/grab v2.0.0+incompatible
99
github.com/dustin/go-humanize v1.0.0

‎go.sum

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@ github.com/VirusTotal/vt-go v0.0.0-20210528074736-45bbe34cc8ab h1:96tkQLYmgypA3W
1919
github.com/VirusTotal/vt-go v0.0.0-20210528074736-45bbe34cc8ab/go.mod h1:u1+HeRyl/gQs67eDgVEWNE7+x+zCyXhdtNVrRJR5YPE=
2020
github.com/VirusTotal/vt-go v0.0.0-20211116094520-07a92e6467b7 h1:1oIDITWcezvHQXrm2RDiYsVpCeRbvxY4MLBAFdoaPWQ=
2121
github.com/VirusTotal/vt-go v0.0.0-20211116094520-07a92e6467b7/go.mod h1:u1+HeRyl/gQs67eDgVEWNE7+x+zCyXhdtNVrRJR5YPE=
22+
github.com/VirusTotal/vt-go v0.0.0-20211209151516-855a1e790678 h1:IVvLDz0INo1rn7wG4OQub9MkyNNaBp0sQgU1apho/mk=
23+
github.com/VirusTotal/vt-go v0.0.0-20211209151516-855a1e790678/go.mod h1:u1+HeRyl/gQs67eDgVEWNE7+x+zCyXhdtNVrRJR5YPE=
2224
github.com/alecthomas/template v0.0.0-20160405071501-a0175ee3bccc/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
2325
github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
2426
github.com/armon/circbuf v0.0.0-20150827004946-bbbad097214e/go.mod h1:3U/XgcO3hCbHZ8TKRvWD2dDTCfh9M9ya+I9JpbB7O8o=

‎utils/client.go

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,11 @@ import (
1717
"container/heap"
1818
"errors"
1919
"fmt"
20-
vt "github.com/VirusTotal/vt-go"
21-
"github.com/spf13/viper"
2220
"os"
2321
"sync"
22+
23+
vt "github.com/VirusTotal/vt-go"
24+
"github.com/spf13/viper"
2425
)
2526

2627
// APIClient represents a VirusTotal API client.

0 commit comments

Comments
 (0)