Add: One-click functionality to check the availability of the current cache

Author: Driver-C

README.md

@@ -37,7 +37,6 @@ Rankings do not differentiate priority levels.
 2. Completing unit test code
 3. Security-related features, such as encrypting configuration files, hiding sensitive information from plain text display, and switching to interactive password input
 4. Support for key-based authentication
-5. One-click functionality to check the availability of the current cache
 
 ## Quick Start
 

README_zh.md

@@ -37,7 +37,6 @@
 2. 完成单元测试代码
 3. 安全相关功能，配置文件加密、隐藏明文显示的敏感信息、密码输入应改为交互式等
 4. 支持秘钥登陆
-5. 一键检查当前缓存是否可用功能
 
 ## 快速开始
 

cmd/cmd.go

@@ -5,6 +5,7 @@ import (
 	"github.com/Driver-C/tryssh/cmd/create"
 	"github.com/Driver-C/tryssh/cmd/delete"
 	"github.com/Driver-C/tryssh/cmd/get"
+	"github.com/Driver-C/tryssh/cmd/prune"
 	"github.com/Driver-C/tryssh/cmd/scp"
 	"github.com/Driver-C/tryssh/cmd/ssh"
 	"github.com/Driver-C/tryssh/cmd/version"
@@ -24,6 +25,7 @@ func NewTrysshCommand() *cobra.Command {
 	rootCmd.AddCommand(create.NewCreateCommand())
 	rootCmd.AddCommand(delete.NewDeleteCommand())
 	rootCmd.AddCommand(get.NewGetCommand())
+	rootCmd.AddCommand(prune.NewPruneCommand())
 	rootCmd.CompletionOptions.DisableDefaultCmd = true
 	return rootCmd
 }

cmd/prune/prune.go

@@ -0,0 +1,37 @@
+package prune
+
+import (
+	"github.com/Driver-C/tryssh/config"
+	"github.com/Driver-C/tryssh/control/prune"
+	"github.com/spf13/cobra"
+	"time"
+)
+
+const (
+	concurrency = 8
+	sshTimeout  = 2 * time.Second
+)
+
+func NewPruneCommand() *cobra.Command {
+	pruneCmd := &cobra.Command{
+		Use:   "prune",
+		Short: "Check if all current caches are available and clear the ones that are not available",
+		Long:  "Check if all current caches are available and clear the ones that are not available",
+		Run: func(cmd *cobra.Command, args []string) {
+			auto, _ := cmd.Flags().GetBool("auto")
+			concurrencyOpt, _ := cmd.Flags().GetInt("concurrency")
+			timeout, _ := cmd.Flags().GetDuration("timeout")
+			configuration := config.LoadConfig()
+			pruneControl := prune.NewPruneController(configuration, auto, timeout, concurrencyOpt)
+			pruneControl.PruneCaches()
+		},
+	}
+	pruneCmd.Flags().BoolP(
+		"auto", "a", false, "Automatically perform concurrent cache optimization without"+
+			" asking for confirmation to delete")
+	pruneCmd.Flags().IntP(
+		"concurrency", "c", concurrency, "Number of multiple requests to perform at a time")
+	pruneCmd.Flags().DurationP("timeout", "t", sshTimeout,
+		"SSH timeout when attempting to log in. It can be \"1s\" or \"1m\" or other duration")
+	return pruneCmd
+}

control/prune/prune.go

@@ -0,0 +1,121 @@
+package prune
+
+import (
+	"bufio"
+	"fmt"
+	"github.com/Driver-C/tryssh/config"
+	"github.com/Driver-C/tryssh/launcher"
+	"github.com/Driver-C/tryssh/launcher/ssh"
+	"github.com/Driver-C/tryssh/utils"
+	"os"
+	"strings"
+	"sync"
+	"time"
+)
+
+type Controller struct {
+	configuration *config.MainConfig
+	auto          bool
+	sshTimeout    time.Duration
+	concurrency   int
+}
+
+func (pc *Controller) PruneCaches() {
+	newServerList := make([]config.ServerListConfig, 0)
+	if pc.auto {
+		newServerList = pc.concurrencyDeleteCache()
+	} else {
+		for _, server := range pc.configuration.ServerLists {
+			lan := &ssh.Launcher{SshConnector: *launcher.GetSshConnectorFromConfig(&server)}
+			// Set timeout
+			lan.SshTimeout = pc.sshTimeout
+			// Determine if connection is possible
+			if err := lan.TryToConnect(); err != nil {
+				if !pc.interactiveDeleteCache(server) {
+					newServerList = append(newServerList, server)
+				}
+			} else {
+				utils.Logger.Infof("Cache %v is still available.", server)
+				newServerList = append(newServerList, server)
+			}
+		}
+	}
+	pc.configuration.ServerLists = newServerList
+	if config.UpdateConfig(pc.configuration) {
+		utils.Logger.Infoln("Update config successful.")
+	} else {
+		utils.Logger.Errorln("Update config failed.")
+	}
+}
+
+func (pc *Controller) interactiveDeleteCache(server config.ServerListConfig) bool {
+	reader := bufio.NewReader(os.Stdin)
+	for {
+		fmt.Printf("Are you sure you want to delete this cache? "+
+			"Please enter \"yes\" to confirm deletion, or \"no\" to cancel. %s\n"+
+			"(yes/no): ", server)
+		stdin, _ := reader.ReadString('\n')
+		// Delete space
+		stdin = strings.TrimSpace(stdin)
+		switch stdin {
+		case "yes":
+			utils.Logger.Infof("The cache %v has been marked for deletion.", server)
+			return true
+		case "no":
+			utils.Logger.Infof("Cache %v skipped.", server)
+			return false
+		default:
+			utils.Logger.Errorln("Input error:", stdin)
+		}
+	}
+}
+
+func (pc *Controller) concurrencyDeleteCache() []config.ServerListConfig {
+	newServerList := make([]config.ServerListConfig, 0)
+	launchersChan := make(chan *ssh.Launcher)
+	var mutex sync.Mutex
+	var wg sync.WaitGroup
+
+	go func(launchersChan chan<- *ssh.Launcher) {
+		for _, server := range pc.configuration.ServerLists {
+			lan := &ssh.Launcher{SshConnector: *launcher.GetSshConnectorFromConfig(&server)}
+			launchersChan <- lan
+		}
+		close(launchersChan)
+	}(launchersChan)
+
+	for i := 0; i < pc.concurrency; i++ {
+		wg.Add(1)
+		go func(launchersChan <-chan *ssh.Launcher, wg *sync.WaitGroup) {
+			defer wg.Done()
+			for {
+				launcherP, ok := <-launchersChan
+				if !ok {
+					break
+				}
+				server := launcher.GetConfigFromSshConnector(&launcherP.SshConnector)
+				launcherP.SshTimeout = pc.sshTimeout
+				if err := launcherP.TryToConnect(); err == nil {
+					utils.Logger.Infof("Cache %v is still available.", server)
+					mutex.Lock()
+					newServerList = append(newServerList, *server)
+					mutex.Unlock()
+				} else {
+					utils.Logger.Infof("The cache %v has been marked for deletion.", server)
+				}
+			}
+		}(launchersChan, &wg)
+	}
+	wg.Wait()
+	return newServerList
+}
+
+func NewPruneController(configuration *config.MainConfig, auto bool, timeout time.Duration,
+	concurrency int) *Controller {
+	return &Controller{
+		configuration: configuration,
+		auto:          auto,
+		sshTimeout:    timeout,
+		concurrency:   concurrency,
+	}
+}