From 5af35c3b3991a9b7d73253944b8f060fdfcf76e7 Mon Sep 17 00:00:00 2001
From: "Victor M. Alvarez" <vmalvarez@virustotal.com>
Date: Tue, 30 Oct 2018 19:30:06 +0100
Subject: [PATCH] Rename "purge" command to "delete", and implement deletion of
 individual notifications and all notifications with a given tag.

---
 cmd/hunting.go | 52 ++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 38 insertions(+), 14 deletions(-)

diff --git a/cmd/hunting.go b/cmd/hunting.go
index 7a8549b..2cd8c20 100644
--- a/cmd/hunting.go
+++ b/cmd/hunting.go
@@ -14,6 +14,7 @@
 package cmd
 
 import (
+	"errors"
 	"fmt"
 	"os"
 	"strconv"
@@ -21,31 +22,56 @@ import (
 
 	"github.com/VirusTotal/vt-go/vt"
 	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
 )
 
-var notificationsPurgeCmdHelp = `Delete all hunting notifications.
+var notificationsDeleteCmdHelp = `Delete hunting notifications.
 
-This command deletes all the malware hunting notifications associated to the
+This command deletes the malware hunting notifications associated to the
 currently configured API key.`
 
-// NewHuntingNotificationsPurgeCmd returns a command for deleting all hunting
+// NewHuntingNotificationsDeleteCmd returns a command for deleting all hunting
 // notifications for the current user.
-func NewHuntingNotificationsPurgeCmd() *cobra.Command {
+func NewHuntingNotificationsDeleteCmd() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:   "purge",
-		Short: "Delete all hunting notifications",
-		Long:  notificationsPurgeCmdHelp,
-
+		Use:   "delete [notification id]...",
+		Short: "Delete hunting notifications",
+		Long:  notificationsDeleteCmdHelp,
 		RunE: func(cmd *cobra.Command, args []string) error {
+			deleteAll := viper.GetBool("all")
+			deleteTag := viper.GetString("with-tag")
+			if len(args) == 0 && !deleteAll && deleteTag == "" {
+				return errors.New("Specify notification id or use --all or --with-tag")
+			}
 			client, err := NewAPIClient()
 			if err != nil {
 				return err
 			}
-			_, err = client.Delete(vt.URL("intelligence/hunting_notifications"))
-			return err
+			if deleteAll {
+				_, err = client.Delete(vt.URL("intelligence/hunting_notifications"))
+			} else if deleteTag != "" {
+				_, err = client.Delete(vt.URL("intelligence/hunting_notifications?tag=%s", deleteTag))
+			} else {
+				var wg sync.WaitGroup
+				for _, arg := range args {
+					wg.Add(1)
+					go func(notificationID string) {
+						url := vt.URL("intelligence/hunting_notifications/%s", notificationID)
+						if _, err := client.Delete(url); err != nil {
+							fmt.Fprintf(os.Stderr, "%v\n", err)
+						}
+						wg.Done()
+					}(arg)
+				}
+				wg.Wait()
+			}
+			return nil
 		},
 	}
 
+	cmd.Flags().BoolP("all", "a", false, "delete all notifications")
+	cmd.Flags().StringP("with-tag", "t", "", "delete notifications with a given tag")
+
 	return cmd
 }
 
@@ -76,7 +102,7 @@ func NewHuntingNotificationsListCmd() *cobra.Command {
 	addLimitFlag(cmd.Flags())
 	addCursorFlag(cmd.Flags())
 
-	cmd.AddCommand(NewHuntingNotificationsPurgeCmd())
+	cmd.AddCommand(NewHuntingNotificationsDeleteCmd())
 
 	return cmd
 }
@@ -90,7 +116,7 @@ func NewHuntingNotificationsCmd() *cobra.Command {
 	}
 
 	cmd.AddCommand(NewHuntingNotificationsListCmd())
-	cmd.AddCommand(NewHuntingNotificationsPurgeCmd())
+	cmd.AddCommand(NewHuntingNotificationsDeleteCmd())
 
 	return cmd
 }
@@ -216,7 +242,6 @@ func NewHuntingRulesetsDeleteCmd() *cobra.Command {
 			if err != nil {
 				return err
 			}
-
 			var wg sync.WaitGroup
 			for _, arg := range args {
 				wg.Add(1)
@@ -228,7 +253,6 @@ func NewHuntingRulesetsDeleteCmd() *cobra.Command {
 					wg.Done()
 				}(arg)
 			}
-
 			wg.Wait()
 			return nil
 		},