From 8666a21d0b1828a21384d5e46a046aeb47e7c870 Mon Sep 17 00:00:00 2001
From: AlexSSD7 <alexandersadovskyi7@protonmail.com>
Date: Wed, 30 Aug 2023 09:43:41 +0100
Subject: [PATCH] Clean cmd + VM image pruning

---
 cmd/clean.go       | 41 +++++++++++++++++++++++++++++++++++
 cmd/root.go        |  1 +
 cmd/utils.go       | 10 ++++++---
 storage/storage.go | 53 +++++++++++++++++++++++++++++++++++++++++-----
 4 files changed, 97 insertions(+), 8 deletions(-)
 create mode 100644 cmd/clean.go

diff --git a/cmd/clean.go b/cmd/clean.go
new file mode 100644
index 0000000..45c85dd
--- /dev/null
+++ b/cmd/clean.go
@@ -0,0 +1,41 @@
+package cmd
+
+import (
+	"bufio"
+	"fmt"
+	"log/slog"
+	"os"
+	"strings"
+
+	"github.com/spf13/cobra"
+)
+
+var cleanCmd = &cobra.Command{
+	Use:   "clean",
+	Short: "Remove all downloaded VM images.",
+	Run: func(cmd *cobra.Command, args []string) {
+		store := createStore()
+
+		fmt.Fprintf(os.Stderr, "Will delete all VM images in the data directory. Proceed? (y/n) > ")
+
+		reader := bufio.NewReader(os.Stdin)
+		answer, err := reader.ReadBytes('\n')
+		if err != nil {
+			slog.Error("Failed to read answer", "error", err.Error())
+			os.Exit(1)
+		}
+
+		if strings.ToLower(string(answer)) != "y\n" {
+			fmt.Fprintf(os.Stderr, "Aborted.\n")
+			os.Exit(2)
+		}
+
+		deleted, err := store.CleanImages(false)
+		if err != nil {
+			slog.Error("Failed to clean images", "error", err.Error())
+			os.Exit(1)
+		}
+
+		slog.Info("Successful VM image cleanup", "deleted", deleted)
+	},
+}
diff --git a/cmd/root.go b/cmd/root.go
index d20958d..e6c0626 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -41,6 +41,7 @@ func init() {
 	rootCmd.AddCommand(lsCmd)
 	rootCmd.AddCommand(runCmd)
 	rootCmd.AddCommand(shellCmd)
+	rootCmd.AddCommand(cleanCmd)
 
 	rootCmd.PersistentFlags().BoolVar(&vmDebugFlag, "vmdebug", false, "Enables the VM debug mode. This will open an accessible VM monitor. You can log in with root user and no password.")
 	rootCmd.PersistentFlags().BoolVar(&unrestrictedNetworkingFlag, "unrestricted-networking", false, "Enables unrestricted networking. This will allow the VM to connect to the internet.")
diff --git a/cmd/utils.go b/cmd/utils.go
index b0ea97b..36498eb 100644
--- a/cmd/utils.go
+++ b/cmd/utils.go
@@ -54,14 +54,19 @@ func doUSBRootCheck() {
 	}
 }
 
-func runVM(passthroughArg string, fn func(context.Context, *vm.VM, *vm.FileManager) int, forwardPortsRules []vm.PortForwardingRule, unrestrictedNetworking bool) int {
+func createStore() *storage.Storage {
 	store, err := storage.NewStorage(slog.With("caller", "storage"), dataDirFlag)
 	if err != nil {
 		slog.Error("Failed to create Linsk data storage", "error", err.Error(), "data-dir", dataDirFlag)
 		os.Exit(1)
 	}
 
-	_, err = store.ValidateImageHashOrDownload()
+	return store
+}
+
+func runVM(passthroughArg string, fn func(context.Context, *vm.VM, *vm.FileManager) int, forwardPortsRules []vm.PortForwardingRule, unrestrictedNetworking bool) int {
+	store := createStore()
+	_, err := store.ValidateImageHashOrDownload()
 	if err != nil {
 		slog.Error("Failed to validate image hash or download image", "error", err.Error())
 		os.Exit(1)
@@ -92,7 +97,6 @@ func runVM(passthroughArg string, fn func(context.Context, *vm.VM, *vm.FileManag
 		ShowDisplay:            vmDebugFlag,
 	}
 
-	// TODO: Alpine image should be downloaded from somewhere.
 	vi, err := vm.NewVM(slog.Default().With("caller", "vm"), vmCfg)
 	if err != nil {
 		slog.Error("Failed to create vm instance", "error", err.Error())
diff --git a/storage/storage.go b/storage/storage.go
index 5b53842..740777e 100644
--- a/storage/storage.go
+++ b/storage/storage.go
@@ -11,6 +11,7 @@ import (
 	"net/http"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"github.com/dustin/go-humanize"
 	"github.com/pkg/errors"
@@ -41,6 +42,35 @@ func NewStorage(logger *slog.Logger, dataDir string) (*Storage, error) {
 	}, nil
 }
 
+func (s *Storage) CleanImages(retainLatest bool) (int, error) {
+	dirEntries, err := os.ReadDir(s.path)
+	if err != nil {
+		return 0, errors.Wrap(err, "read dir")
+	}
+
+	localImagePath := s.GetLocalImagePath()
+
+	var deleted int
+
+	for _, entry := range dirEntries {
+		fullPath := filepath.Join(s.path, entry.Name())
+		if !strings.HasSuffix(fullPath, ".qcow2") {
+			continue
+		}
+
+		if !(retainLatest && fullPath == localImagePath) {
+			s.logger.Warn("Removing VM image", "path", fullPath)
+			err = os.Remove(fullPath)
+			if err != nil {
+				return deleted, errors.Wrapf(err, "remove vm image (path '%v')", fullPath)
+			}
+			deleted++
+		}
+	}
+
+	return deleted, nil
+}
+
 func (s *Storage) GetLocalImagePath() string {
 	return filepath.Join(s.path, hex.EncodeToString(imageHash[:])+".qcow2")
 }
@@ -143,16 +173,29 @@ func (s *Storage) ValidateImageHash() error {
 }
 
 func (s *Storage) ValidateImageHashOrDownload() (bool, error) {
+	var downloaded bool
 	err := s.ValidateImageHash()
-	if err == nil {
-		return false, nil
+	if err != nil {
+		if errors.Is(err, os.ErrNotExist) {
+			err = s.DownloadImage()
+			if err != nil {
+				return false, errors.Wrap(err, "download iamge")
+			}
+
+			downloaded = true
+		} else {
+			return false, errors.Wrap(err, "validate image hash")
+		}
 	}
 
-	if errors.Is(err, os.ErrNotExist) {
-		return true, errors.Wrap(s.DownloadImage(), "download image")
+	deletedImagesCount, err := s.CleanImages(true)
+	if err != nil {
+		s.logger.Warn("Failed to prune old VM images", "error", err, "deleted", deletedImagesCount)
+	} else {
+		s.logger.Info("Pruned old VM images", "deleted", deletedImagesCount)
 	}
 
-	return false, err
+	return downloaded, err
 }
 
 func copyWithProgress(dst io.Writer, src io.Reader, blockSize int, length int64, report func(int, float64)) (int, error) {