feat: support downloading files in async, better TUI, extract flag (#15)

Co-authored-by: Ayoub Faouzi <ayoubfaouzi@users.noreply.github.com>

Author: ayoubfaouzi

.golangci.yaml

@@ -57,8 +57,15 @@ saferwall-cli rescan <sha256>
 
 ### Download
 
-Download files by their SHA256 hash. You can also download a batch of samples from a text file.
+Download a sample by its SHA256 hash, or provide a text file with one hash per line to download in batch.
 
 ```sh
-saferwall-cli download --hash <sha256>
+# Single sample
+saferwall-cli download <sha256>
+
+# Batch from a text file
+saferwall-cli download hashes.txt
+
+# Extract from zip (password: infected) instead of keeping the .zip
+saferwall-cli download -x <sha256>
 ```

README.md

@@ -5,88 +5,87 @@
 package cmd
 
 import (
-	"bytes"
+	"fmt"
 	"log"
 	"os"
 	"path/filepath"
 	"strings"
 
+	tea "github.com/charmbracelet/bubbletea"
 	"github.com/saferwall/cli/internal/util"
 	"github.com/saferwall/cli/internal/webapi"
 	"github.com/spf13/cobra"
 )
 
-var sha256Flag string
-var txtFlag string
 var outputFlag string
+var extractFlag bool
 
 func init() {
 	ex, err := os.Executable()
 	if err != nil {
 		panic(err)
 	}
 
-	downloadCmd.Flags().StringVarP(&sha256Flag, "hash", "s", "", "SHA256 hash to download")
-	downloadCmd.Flags().StringVarP(&txtFlag, "txt", "t", "", "Download all hashes in a text file, separate by a line break")
 	downloadCmd.Flags().StringVarP(&outputFlag, "output", "o", filepath.Dir(ex),
 		"Destination directory where to save samples. (default=current dir)")
+	downloadCmd.Flags().IntVarP(&parallelFlag, "parallel", "p", 4,
+		"Number of files to download in parallel")
+	downloadCmd.Flags().BoolVarP(&extractFlag, "extract", "x", false,
+		"Extract samples from zip (password: infected)")
 }
 
 var downloadCmd = &cobra.Command{
-	Use:   "download",
-	Short: "Download a sample(s)",
-	Long:  `Download a binary sample given a sha256`,
+	Use:   "download <sha256|file.txt>",
+	Short: "Download a sample (and its artifacts)",
+	Long:  `Download a binary sample given a SHA256 hash, or a batch of samples from a text file containing one hash per line.`,
+	Args:  cobra.ExactArgs(1),
 	Run: func(cmd *cobra.Command, args []string) {
+		arg := args[0]
 
-		// Login to saferwall web service
+		// Login to saferwall web service.
 		webSvc := webapi.New(cfg.Credentials.URL)
 		token, err := webSvc.Login(cfg.Credentials.Username, cfg.Credentials.Password)
 		if err != nil {
 			log.Fatalf("failed to login to saferwall web service")
 		}
 
-		// download a single binary.
-		if sha256Flag != "" {
-			download(sha256Flag, token, webSvc)
-		} else if txtFlag != "" {
-			// Download a list of sha256 hashes.
-			data, err := util.ReadAll(txtFlag)
-			if err != nil {
-				log.Fatalf("failed to read to SHA256 hashes from txt file: %v", txtFlag)
-			}
-
-			sha256list := strings.Split(string(data), "\n")
-			for _, sha256 := range sha256list {
-				if len(sha256) >= 64 {
-					err = download(sha256, token, webSvc)
-					if err != nil {
-						log.Fatalf("failed to download sample (%s): %v", sha256, err)
-					}
-				}
-			}
+		hashes := collectHashes(arg)
+		if len(hashes) == 0 {
+			log.Fatalf("no valid SHA256 hashes found in %q", arg)
 		}
+
+		downloadFiles(webSvc, token, hashes)
 	},
 }
 
-func download(sha256, token string, web webapi.Service) error {
-	var err error
-	var data bytes.Buffer
-	var destPath string
+// collectHashes returns a list of SHA256 hashes from the argument.
+// If arg is a SHA256 hash, it returns a single-element slice.
+// Otherwise it treats arg as a file path and reads hashes from it.
+func collectHashes(arg string) []string {
+	if sha256Re.MatchString(arg) {
+		return []string{arg}
+	}
 
-	log.Printf("downloading %s to %s", sha256, outputFlag)
-	dataContent, err := web.Download(sha256, token)
+	data, err := util.ReadAll(arg)
 	if err != nil {
-		log.Fatalf("failed to download %s, err: %v", sha256, err)
-		return err
+		log.Fatalf("failed to read SHA256 hashes from file: %s", arg)
 	}
-	data = *dataContent
 
-	filename := sha256 + ".zip"
-	destPath = filepath.Join(outputFlag, filename)
-	_, err = util.WriteBytesFile(destPath, &data)
-	if err != nil {
-		return err
+	var hashes []string
+	for _, line := range strings.Split(string(data), "\n") {
+		line = strings.TrimSpace(line)
+		if sha256Re.MatchString(line) {
+			hashes = append(hashes, line)
+		}
 	}
+	return hashes
+}
 
-	return nil
+func downloadFiles(web webapi.Service, token string, hashes []string) {
+	model := newDownloadModel(hashes, web, token, outputFlag, parallelFlag, extractFlag)
+	p := tea.NewProgram(model)
+	if _, err := p.Run(); err != nil {
+		fmt.Fprintf(os.Stderr, "TUI error: %v\n", err)
+		os.Exit(1)
+	}
 }