feat: make scan command look pretty (#12)

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

Author: ayoubfaouzi

.pre-commit-config.yaml

@@ -1,24 +1,8 @@
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.2.0
+    rev: v6.0.0
     hooks:
       - id: trailing-whitespace
       - id: end-of-file-fixer
       - id: check-yaml
       - id: check-added-large-files
-  - repo: https://github.com/dnephin/pre-commit-golang
-    rev: master
-    hooks:
-      - id: go-fmt
-      - id: go-vet
-      - id: go-lint
-      - id: go-imports
-      - id: go-cyclo
-        args: [-over=15]
-      - id: validate-toml
-      - id: no-go-testing
-      - id: golangci-lint
-      - id: go-critic
-      - id: go-unit-tests
-      - id: go-build
-      - id: go-mod-tidy

cmd/rescan.go

@@ -6,27 +6,27 @@ package cmd
 
 import (
 	"log"
-	"runtime"
 	"strings"
+	"sync"
 	"time"
 
-	"github.com/gammazero/workerpool"
 	"github.com/saferwall/cli/internal/util"
 	"github.com/saferwall/cli/internal/webapi"
 	"github.com/spf13/cobra"
 )
 
 var (
-	fileHash string
+	rescanFilePath string
+	fileHash       string
 )
 
 func init() {
-	reScanCmd.Flags().StringVarP(&filePath, "path", "p", "",
+	reScanCmd.Flags().StringVarP(&rescanFilePath, "path", "p", "",
 		"File name or path containing list of SHA256 to scan")
 	reScanCmd.Flags().StringVarP(&fileHash, "hash", "s", "",
 		"SHA256 of the file to rescan")
-	reScanCmd.Flags().BoolVarP(&asyncScanFlag, "async", "a", false,
-		"Scan files in parallel")
+	reScanCmd.Flags().IntVar(&parallelFlag, "parallel", 1,
+		"Number of files to rescan in parallel")
 	reScanCmd.Flags().BoolVarP(&enableDetonationFlag, "enableDetonation", "d", false,
 		"Skip detonation")
 	reScanCmd.Flags().IntVarP(&timeoutFlag, "timeout", "t", 15,
@@ -37,42 +37,23 @@ func init() {
 
 // reScanFile re-scans a list of SHA256.
 func reScanFile(web webapi.Service, shaList []string, token string) error {
+	sem := make(chan struct{}, parallelFlag)
+	var wg sync.WaitGroup
 
-	if asyncScanFlag {
-		// Create a worker pool
-		maxWorkers := runtime.GOMAXPROCS(0)
-		wp := workerpool.New(maxWorkers)
-
-		for _, sha256 := range shaList {
-			wp.Submit(func() {
-				log.Printf("rescanning %s", sha256)
-				err := web.Rescan(sha256, token, osFlag, enableDetonationFlag, timeoutFlag)
-				if err != nil {
-					log.Fatalf("failed to rescan file: %v", sha256)
-				}
-
-				time.Sleep(2 * time.Second)
-			})
-		}
-		wp.StopWait()
-		return nil
-	}
-
-	// Sequentially scan the files.
 	for _, sha256 := range shaList {
-
-		log.Printf("re-scanning %s", sha256)
-		err := web.Rescan(sha256, token, osFlag, enableDetonationFlag, timeoutFlag)
-		if err != nil {
-			log.Fatalf("failed to rescan file: %v", sha256)
-		}
-
-		if len(shaList) > 1 {
-			time.Sleep(10 * time.Second)
-		}
-
+		sem <- struct{}{}
+		wg.Add(1)
+		go func() {
+			defer func() { <-sem; wg.Done() }()
+			log.Printf("rescanning %s", sha256)
+			err := web.Rescan(sha256, token, osFlag, enableDetonationFlag, timeoutFlag)
+			if err != nil {
+				log.Printf("failed to rescan file: %v", sha256)
+			}
+			time.Sleep(2 * time.Second)
+		}()
 	}
-
+	wg.Wait()
 	return nil
 }
 
@@ -91,8 +72,8 @@ var reScanCmd = &cobra.Command{
 
 		// Read the txt file containing the list of hashes to rescan.
 		var sha256List []string
-		if filePath != "" {
-			data, err := util.ReadAll(filePath)
+		if rescanFilePath != "" {
+			data, err := util.ReadAll(rescanFilePath)
 			if err != nil {
 				log.Fatalf("failed to read txt file")
 			}

cmd/root.go

@@ -5,7 +5,6 @@
 package cmd
 
 import (
-	"fmt"
 	"log"
 	"path/filepath"
 
@@ -14,10 +13,6 @@ import (
 	"github.com/spf13/cobra"
 )
 
-const (
-	version = "0.5.0"
-)
-
 var cfg config.Config
 
 var rootCmd = &cobra.Command{
@@ -38,16 +33,7 @@ upload, scan samples from your drive, or download samples.
 For more details see the github repo at https://github.com/saferwall
 `,
 	Run: func(cmd *cobra.Command, args []string) {
-		fmt.Printf("You are using version %s\n", version)
-	},
-}
-
-var versionCmd = &cobra.Command{
-	Use:   "version",
-	Short: "Print the version number",
-	Long:  "Print the version number",
-	Run: func(cmd *cobra.Command, args []string) {
-		fmt.Printf("You are using version %s\n", version)
+		cmd.Help()
 	},
 }
 

cmd/scan.go

@@ -5,46 +5,85 @@
 package cmd
 
 import (
+	"fmt"
 	"log"
 	"os"
 	"path/filepath"
-	"runtime"
 	"time"
 
-	"github.com/gammazero/workerpool"
-	"github.com/saferwall/cli/internal/util"
+	tea "github.com/charmbracelet/bubbletea"
+	"github.com/saferwall/cli/internal/entity"
 	"github.com/saferwall/cli/internal/webapi"
 	"github.com/spf13/cobra"
 )
 
+const (
+	statusQueued    = 1
+	statusScanning  = 2
+	statusCompleted = 3
+
+	pollInterval = 5 * time.Second
+)
+
 // Used for flags.
-var filePath string
 var forceRescanFlag bool
-var asyncScanFlag bool
+var parallelFlag int
 var enableDetonationFlag bool
 var timeoutFlag int
 var osFlag string
 
 func init() {
-	scanCmd.Flags().StringVarP(&filePath, "path", "p", "",
-		"File name or path to scan (required)")
 	scanCmd.Flags().BoolVarP(&forceRescanFlag, "force", "f", false,
 		"Force rescan the file if it exists")
-	scanCmd.Flags().BoolVarP(&asyncScanFlag, "async", "a", false,
-		"Scan files in parallel")
+	scanCmd.Flags().IntVarP(&parallelFlag, "parallel", "p", 1,
+		"Number of files to scan in parallel")
 	scanCmd.Flags().BoolVarP(&enableDetonationFlag, "enableDetonation", "d", false,
 		"Skip detonation")
 	scanCmd.Flags().IntVarP(&timeoutFlag, "timeout", "t", 15,
 		"Detonation duration in seconds")
 	scanCmd.Flags().StringVarP(&osFlag, "os", "o", "win-10",
 		"Preferred OS for detonation, choice(win-7 | win-10)")
-	scanCmd.MarkFlagRequired("path")
+}
+
+type scanSummary struct {
+	SHA256         string     `json:"sha256"`
+	Classification string     `json:"classification"`
+	FileFormat     string     `json:"file_format"`
+	FileExtension  string     `json:"file_extension"`
+	MultiAV        *avSummary `json:"multiav,omitempty"`
+}
+
+type avSummary struct {
+	Positives    int `json:"positives"`
+	EnginesCount int `json:"engines_count"`
+}
 
+func buildScanSummary(file entity.File) scanSummary {
+	s := scanSummary{
+		SHA256:         file.SHA256,
+		Classification: file.Classification,
+		FileFormat:     file.Format,
+		FileExtension:  file.Extension,
+	}
+
+	if lastScan, ok := file.MultiAV["last_scan"].(map[string]any); ok {
+		if stats, ok := lastScan["stats"].(map[string]any); ok {
+			av := &avSummary{}
+			if v, ok := stats["positives"].(float64); ok {
+				av.Positives = int(v)
+			}
+			if v, ok := stats["engines_count"].(float64); ok {
+				av.EnginesCount = int(v)
+			}
+			s.MultiAV = av
+		}
+	}
+
+	return s
 }
 
 // scanFile scans an individual file or a directory.
 func scanFile(web webapi.Service, filePath, token string) error {
-
 	_, err := os.Stat(filePath)
 	if os.IsNotExist(err) {
 		log.Printf("file path [%s] does not exists", filePath)
@@ -60,98 +99,20 @@ func scanFile(web webapi.Service, filePath, token string) error {
 		return nil
 	})
 
-	if asyncScanFlag {
-
-		// Create a worker pool
-		maxWorkers := runtime.GOMAXPROCS(0)
-		wp := workerpool.New(maxWorkers)
-
-		// Upload files
-		for _, filename := range fileList {
-			filename := filename
-			wp.Submit(func() {
-
-				// Get sha256
-				data, err := os.ReadFile(filename)
-				if err != nil {
-					log.Fatalf("failed to read file: %v", filename)
-				}
-				sha256 := util.GetSha256(data)
-
-				// Check if we the file exists in the DB.
-				exists, err := web.FileExists(sha256)
-				if err != nil {
-					log.Fatalf("failed to check existence of file: %v", filename)
-				}
-
-				// Upload the file to be scanned, this will automatically trigger a scan request.
-				if !exists {
-					_, err = web.Scan(filename, token, osFlag, enableDetonationFlag, timeoutFlag)
-					if err != nil {
-						log.Fatalf("failed to upload file: %v", filename)
-					}
-				} else {
-					// Force rescan the file
-					if forceRescanFlag {
-						err = web.Rescan(sha256, token, osFlag, enableDetonationFlag, timeoutFlag)
-						if err != nil {
-							log.Fatalf("failed to rescan file: %v", filename)
-						}
-					}
-				}
-
-				time.Sleep(2 * time.Second)
-			})
-		}
-		wp.StopWait()
-		return nil
-	}
-
-	// Sequentially scan the files.
-	for _, filename := range fileList {
-		// Get sha256
-		data, err := os.ReadFile(filename)
-		if err != nil {
-			log.Fatalf("failed to read file: %v", filename)
-		}
-		sha256 := util.GetSha256(data)
-
-		log.Printf("processing %s", sha256)
-
-		// Check if we the file exists in the DB.
-		exists, err := web.FileExists(sha256)
-		if err != nil {
-			log.Fatalf("failed to check existence of file: %s, error: %v", filename, err)
-		}
-
-		// Upload the file to be scanned, this will automatically
-		// trigger a scan request.
-		if !exists {
-			body, err := web.Scan(filename, token, osFlag, enableDetonationFlag, timeoutFlag)
-			if err != nil {
-				log.Fatalf("failed to upload file: %s, error: %v", filename, err)
-			}
-			log.Print(body)
-			time.Sleep(10 * time.Second)
-		} else {
-			// Force re-scan the file
-			if forceRescanFlag {
-				err = web.Rescan(sha256, token, osFlag, enableDetonationFlag, timeoutFlag)
-				if err != nil {
-					log.Fatalf("failed to re-scan file: %v", filename)
-				}
-			}
-		}
-
+	// Launch TUI scan with the configured parallelism.
+	model := newScanModel(fileList, web, token, parallelFlag)
+	p := tea.NewProgram(model)
+	if _, err := p.Run(); err != nil {
+		return fmt.Errorf("TUI error: %w", err)
 	}
-
 	return nil
 }
 
 var scanCmd = &cobra.Command{
-	Use:   "scan",
+	Use:   "scan <path>",
 	Short: "Submit a scan request of a file using its hash",
 	Long:  `Scans the file`,
+	Args:  cobra.ExactArgs(1),
 	Run: func(cmd *cobra.Command, args []string) {
 
 		// login to saferwall web service
@@ -161,6 +122,6 @@ var scanCmd = &cobra.Command{
 			log.Fatalf("failed to login to saferwall web service")
 		}
 
-		scanFile(webSvc, filePath, token)
+		scanFile(webSvc, args[0], token)
 	},
 }