Skip to content

Commit 963f9cc

Browse files
joseotorojoseotoro
authored
feat: Interact with threat profiles API (#109)
1 parent a903645 commit 963f9cc

File tree

3 files changed

+359
-1
lines changed

3 files changed

+359
-1
lines changed

‎cmd/init.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ func NewInitCmd() *cobra.Command {
5151

5252
Run: func(cmd *cobra.Command, args []string) {
5353

54-
fmt.Printf(vtBanner)
54+
fmt.Print(vtBanner)
5555

5656
apiKey := cmd.Flags().Lookup("apikey").Value.String()
5757

‎cmd/threat_profile.go‎

Lines changed: 357 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,357 @@
1+
package cmd
2+
3+
import (
4+
"fmt"
5+
"os"
6+
7+
"github.com/VirusTotal/vt-go"
8+
"github.com/spf13/cobra"
9+
"github.com/spf13/viper"
10+
11+
"github.com/VirusTotal/vt-cli/utils"
12+
)
13+
14+
var threatProfileCmdHelp = `Get information about one or more Threat Profiles.
15+
16+
This command receives one or more Threat Profile IDs and returns information about them.
17+
The information for each profile is returned in the same order as the IDs are passed to the command.
18+
19+
If the command receives a single hyphen (-) the IDs will be read from the standard input, one per line.`
20+
21+
var threatProfileCmdExample = ` vt threatprofile <profile_id_1>
22+
vt threatprofile <profile_id_1> <profile_id_2>
23+
cat list_of_profile_ids | vt threatprofile -`
24+
25+
// NewThreatProfileCmd returns a new instance of the 'threatprofile' command.
26+
func NewThreatProfileCmd() *cobra.Command {
27+
cmd := &cobra.Command{
28+
Use: "threatprofile [id]...",
29+
Short: "Get information about Threat Profiles",
30+
Long: threatProfileCmdHelp,
31+
Example: threatProfileCmdExample,
32+
Args: cobra.MinimumNArgs(1), // For fetching specific profiles by ID
33+
34+
RunE: func(cmd *cobra.Command, args []string) error {
35+
p, err := NewPrinter(cmd)
36+
if err != nil {
37+
return err
38+
}
39+
return p.GetAndPrintObjects(
40+
"threat_profiles/%s",
41+
utils.StringReaderFromCmdArgs(args),
42+
nil) // No specific regexp for ID needed
43+
},
44+
}
45+
46+
addRelationshipCmds(cmd, "threat_profiles", "threat_profile", "[id]")
47+
addThreadsFlag(cmd.Flags())
48+
addIncludeExcludeFlags(cmd.Flags())
49+
addIDOnlyFlag(cmd.Flags())
50+
51+
cmd.AddCommand(NewThreatProfileListCmd())
52+
cmd.AddCommand(NewThreatProfileCreateCmd())
53+
cmd.AddCommand(NewThreatProfileUpdateCmd())
54+
cmd.AddCommand(NewThreatProfileDeleteCmd())
55+
56+
return cmd
57+
}
58+
59+
var threatProfileListCmdHelp = `List Threat Profiles.`
60+
var threatProfileListCmdExample = ` vt threatprofile list
61+
vt threatprofile list --filter "name:APT" --limit 10
62+
vt threatprofile list --cursor <cursor_value>`
63+
64+
// NewThreatProfileListCmd returns a command for listing Threat Profiles.
65+
func NewThreatProfileListCmd() *cobra.Command {
66+
cmd := &cobra.Command{
67+
Use: "list",
68+
Short: "List Threat Profiles",
69+
Long: threatProfileListCmdHelp,
70+
Example: threatProfileListCmdExample,
71+
RunE: func(cmd *cobra.Command, args []string) error {
72+
p, err := NewPrinter(cmd)
73+
if err != nil {
74+
return err
75+
}
76+
return p.PrintCollection(vt.URL("threat_profiles"))
77+
},
78+
}
79+
80+
addIncludeExcludeFlags(cmd.Flags())
81+
addIDOnlyFlag(cmd.Flags())
82+
addFilterFlag(cmd.Flags())
83+
addLimitFlag(cmd.Flags())
84+
addCursorFlag(cmd.Flags())
85+
86+
return cmd
87+
}
88+
89+
var threatProfileUpdateCmdHelp = `Update a Threat Profile.
90+
91+
This command updates an existing Threat Profile with the specified ID.
92+
You can update attributes like name, interests, and recommendation configuration.`
93+
94+
var threatProfileUpdateCmdExample = ` vt threatprofile update <profile_id> --name "Updated Name"
95+
vt threatprofile update <profile_id> --targeted-region "US,CA" --actor-motivation "cybercrime"`
96+
97+
// NewThreatProfileUpdateCmd returns a command for updating a Threat Profile.
98+
func NewThreatProfileUpdateCmd() *cobra.Command {
99+
cmd := &cobra.Command{
100+
Use: "update [id]",
101+
Short: "Update a Threat Profile",
102+
Long: threatProfileUpdateCmdHelp,
103+
Example: threatProfileUpdateCmdExample,
104+
Args: cobra.ExactArgs(1), // Threat Profile ID is required
105+
106+
RunE: func(cmd *cobra.Command, args []string) error {
107+
client, err := NewAPIClient()
108+
if err != nil {
109+
return err
110+
}
111+
printer, err := NewPrinter(cmd)
112+
if err != nil {
113+
return err
114+
}
115+
116+
profileID := args[0]
117+
threatProfile := vt.NewObjectWithID("threat_profile", profileID)
118+
119+
// Check and set flags for attributes
120+
if cmd.Flags().Changed("name") {
121+
threatProfile.SetString("name", viper.GetString("name"))
122+
}
123+
124+
// Check and set flags for interests
125+
interestsData := make(map[string]interface{})
126+
interestsChanged := false
127+
if cmd.Flags().Changed("targeted-industry") {
128+
interestsData["INTEREST_TYPE_TARGETED_INDUSTRY"] = viper.GetStringSlice("targeted-industry")
129+
interestsChanged = true
130+
}
131+
if cmd.Flags().Changed("targeted-region") {
132+
interestsData["INTEREST_TYPE_TARGETED_REGION"] = viper.GetStringSlice("targeted-region")
133+
interestsChanged = true
134+
}
135+
if cmd.Flags().Changed("source-region") {
136+
interestsData["INTEREST_TYPE_SOURCE_REGION"] = viper.GetStringSlice("source-region")
137+
interestsChanged = true
138+
}
139+
if cmd.Flags().Changed("malware-role") {
140+
interestsData["INTEREST_TYPE_MALWARE_ROLE"] = viper.GetStringSlice("malware-role")
141+
interestsChanged = true
142+
}
143+
if cmd.Flags().Changed("actor-motivation") {
144+
interestsData["INTEREST_TYPE_ACTOR_MOTIVATION"] = viper.GetStringSlice("actor-motivation")
145+
interestsChanged = true
146+
}
147+
if interestsChanged {
148+
threatProfile.Set("interests", interestsData)
149+
}
150+
151+
// Check and set flags for recommendation_config
152+
recommendationConfigData := make(map[string]interface{})
153+
recommendationConfigChanged := false
154+
if cmd.Flags().Changed("max-recs-per-type") {
155+
recommendationConfigData["max_recs_per_type"] = viper.GetInt("max-recs-per-type")
156+
recommendationConfigChanged = true
157+
}
158+
if cmd.Flags().Changed("min-categories-matched") {
159+
recommendationConfigData["min_categories_matched"] = viper.GetInt("min-categories-matched")
160+
recommendationConfigChanged = true
161+
}
162+
if cmd.Flags().Changed("max-days-since-last-seen") {
163+
recommendationConfigData["max_days_since_last_seen"] = viper.GetInt("max-days-since-last-seen")
164+
recommendationConfigChanged = true
165+
}
166+
if recommendationConfigChanged {
167+
threatProfile.Set("recommendation_config", recommendationConfigData)
168+
}
169+
170+
// Need to check if *any* flag was changed besides the default ones
171+
// (like --format, --apikey, etc.). If only the ID is provided
172+
// without any update flags, it should probably error or do nothing.
173+
// Let's check if any of the specific update flags were changed.
174+
updateFlagsChanged := cmd.Flags().Changed("name") ||
175+
interestsChanged ||
176+
recommendationConfigChanged
177+
178+
if !updateFlagsChanged {
179+
return fmt.Errorf("no update flags provided. Use --help for available flags")
180+
}
181+
182+
if err := client.PatchObject(vt.URL("threat_profiles/%s", profileID), threatProfile); err != nil {
183+
return err
184+
}
185+
186+
// Fetch the updated object to print the full details, as PatchObject might not return all attributes
187+
updatedThreatProfile, err := client.GetObject(vt.URL("threat_profiles/%s", profileID))
188+
if err != nil {
189+
// If fetching the updated object fails, at least report the patch was successful
190+
fmt.Fprintf(os.Stderr, "Warning: Failed to fetch updated threat profile details: %v\n", err)
191+
fmt.Printf("Threat profile %s updated successfully.\n", profileID)
192+
return nil
193+
}
194+
195+
if viper.GetBool("identifiers-only") {
196+
fmt.Printf("%s\n", updatedThreatProfile.ID())
197+
} else {
198+
return printer.PrintObject(updatedThreatProfile)
199+
}
200+
201+
return nil
202+
},
203+
}
204+
205+
// Add flags for updatable attributes
206+
cmd.Flags().StringP("name", "n", "", "Threat Profile's name")
207+
208+
// Flags for interests (optional, can be updated)
209+
cmd.Flags().StringSlice("targeted-industry", []string{}, "List of targeted industries (comma-separated)")
210+
cmd.Flags().StringSlice("targeted-region", []string{}, "List of targeted regions (comma-separated)")
211+
cmd.Flags().StringSlice("source-region", []string{}, "List of source regions (comma-separated)")
212+
cmd.Flags().StringSlice("malware-role", []string{}, "List of malware roles (comma-separated)")
213+
cmd.Flags().StringSlice("actor-motivation", []string{}, "List of actors’ motivations (comma-separated)")
214+
215+
// Flags for recommendation_config (optional, can be updated)
216+
cmd.Flags().Int("max-recs-per-type", 0, "Max recommendations per type (1-20)") // Use 0 as default to detect if set
217+
cmd.Flags().Int("min-categories-matched", 0, "Min matching categories for recommendation (1-5)")
218+
cmd.Flags().Int("max-days-since-last-seen", 0, "Max lookback period in days for recommendations (1-365)")
219+
220+
addIncludeExcludeFlags(cmd.Flags())
221+
addIDOnlyFlag(cmd.Flags())
222+
223+
return cmd
224+
}
225+
226+
var threatProfileDeleteCmdHelp = `Delete one or more Threat Profiles.
227+
228+
This command receives one or more Threat Profile IDs and deletes them.
229+
The command will ask for confirmation before deleting.`
230+
231+
var threatProfileDeleteCmdExample = ` vt threatprofile delete <profile_id_1>
232+
vt threatprofile delete <profile_id_1> <profile_id_2>
233+
cat list_of_profile_ids | vt threatprofile delete -`
234+
235+
// NewThreatProfileDeleteCmd returns a command for deleting Threat Profiles.
236+
func NewThreatProfileDeleteCmd() *cobra.Command {
237+
cmd := &cobra.Command{
238+
Use: "delete [id]...",
239+
Short: "Delete Threat Profiles",
240+
Long: threatProfileDeleteCmdHelp,
241+
Example: threatProfileDeleteCmdExample,
242+
Args: cobra.MinimumNArgs(1),
243+
244+
RunE: func(cmd *cobra.Command, args []string) error {
245+
client, err := NewAPIClient()
246+
if err != nil {
247+
return err
248+
}
249+
for _, id := range args {
250+
if _, err := client.Delete(vt.URL("threat_profiles/%s", id)); err != nil {
251+
return err
252+
}
253+
}
254+
return nil
255+
},
256+
}
257+
return cmd
258+
}
259+
260+
var createThreatProfileCmdHelp = `Creates a Threat Profile.
261+
262+
This command creates a new Threat Profile with the specified name, description,
263+
interests, and recommendation configuration.
264+
For interest types, provide comma-separated values if multiple values are needed for a single interest type flag.`
265+
266+
var createThreatProfileCmdExample = ` vt threatprofile create --name "My New Threat Profile" --targeted-region "US,ES"`
267+
268+
// NewThreatProfileCreateCmd returns a command for creating a Threat Profile.
269+
func NewThreatProfileCreateCmd() *cobra.Command {
270+
cmd := &cobra.Command{
271+
Use: "create",
272+
Short: "Create a Threat Profile",
273+
Long: createThreatProfileCmdHelp,
274+
Example: createThreatProfileCmdExample,
275+
RunE: func(cmd *cobra.Command, args []string) error {
276+
client, err := NewAPIClient()
277+
if err != nil {
278+
return err
279+
}
280+
printer, err := NewPrinter(cmd)
281+
if err != nil {
282+
return err
283+
}
284+
285+
threatProfile := vt.NewObject("threat_profile")
286+
threatProfile.SetString("name", viper.GetString("name"))
287+
288+
// Optional interests
289+
interestsData := make(map[string]interface{})
290+
if viper.IsSet("targeted-industry") {
291+
interestsData["INTEREST_TYPE_TARGETED_INDUSTRY"] = viper.GetStringSlice("targeted-industry")
292+
}
293+
if viper.IsSet("targeted-region") {
294+
interestsData["INTEREST_TYPE_TARGETED_REGION"] = viper.GetStringSlice("targeted-region")
295+
}
296+
if viper.IsSet("source-region") {
297+
interestsData["INTEREST_TYPE_SOURCE_REGION"] = viper.GetStringSlice("source-region")
298+
}
299+
if viper.IsSet("malware-role") {
300+
interestsData["INTEREST_TYPE_MALWARE_ROLE"] = viper.GetStringSlice("malware-role")
301+
}
302+
if viper.IsSet("actor-motivation") {
303+
interestsData["INTEREST_TYPE_ACTOR_MOTIVATION"] = viper.GetStringSlice("actor-motivation")
304+
}
305+
306+
if len(interestsData) > 0 {
307+
threatProfile.Set("interests", interestsData)
308+
}
309+
310+
// Optional recommendation_config
311+
recommendationConfigData := make(map[string]interface{})
312+
if viper.IsSet("max-recs-per-type") {
313+
recommendationConfigData["max_recs_per_type"] = viper.GetInt("max-recs-per-type")
314+
}
315+
if viper.IsSet("min-categories-matched") {
316+
recommendationConfigData["min_categories_matched"] = viper.GetInt("min-categories-matched")
317+
}
318+
if viper.IsSet("max-days-since-last-seen") {
319+
recommendationConfigData["max_days_since_last_seen"] = viper.GetInt("max-days-since-last-seen")
320+
}
321+
if len(recommendationConfigData) > 0 {
322+
threatProfile.Set("recommendation_config", recommendationConfigData)
323+
}
324+
325+
if err := client.PostObject(vt.URL("threat_profiles"), threatProfile); err != nil {
326+
return err
327+
}
328+
329+
if viper.GetBool("identifiers-only") {
330+
fmt.Printf("%s\n", threatProfile.ID())
331+
} else {
332+
return printer.PrintObject(threatProfile)
333+
}
334+
return nil
335+
},
336+
}
337+
338+
cmd.Flags().StringP("name", "n", "", "Threat Profile's name (required)")
339+
_ = cmd.MarkFlagRequired("name")
340+
341+
// Flags for interests
342+
cmd.Flags().StringSlice("targeted-industry", []string{}, "List of targeted industries (comma-separated)")
343+
cmd.Flags().StringSlice("targeted-region", []string{}, "List of targeted regions (comma-separated)")
344+
cmd.Flags().StringSlice("source-region", []string{}, "List of source regions (comma-separated)")
345+
cmd.Flags().StringSlice("malware-role", []string{}, "List of malware roles (comma-separated)")
346+
cmd.Flags().StringSlice("actor-motivation", []string{}, "List of actors’ motivations (comma-separated)")
347+
348+
// Flags for recommendation_config
349+
cmd.Flags().Int("max-recs-per-type", 10, "Max recommendations per type (1-20, default 10 if not set by API)") // Default to 0 to check if set
350+
cmd.Flags().Int("min-categories-matched", 1, "Min matching categories for recommendation (1-5, default 1 if not set by API)")
351+
cmd.Flags().Int("max-days-since-last-seen", 180, "Max lookback period in days for recommendations (1-365, default 180 if not set by API)")
352+
353+
addIncludeExcludeFlags(cmd.Flags())
354+
addIDOnlyFlag(cmd.Flags())
355+
356+
return cmd
357+
}

‎cmd/vt.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,7 @@ func NewVTCommand() *cobra.Command {
8686
cmd.AddCommand(NewVersionCmd())
8787
cmd.AddCommand(NewMonitorCmd())
8888
cmd.AddCommand(NewMonitorPartnerCmd())
89+
cmd.AddCommand(NewThreatProfileCmd())
8990

9091
return cmd
9192
}

0 commit comments

Comments
 (0)