Update

.gitignore

@@ -1,7 +1,7 @@
 bin/
 vendor
 ssh-benchmark
-
+benchmark.zip
 ---
 .vscode/*
 !.vscode/settings.json

main.go

@@ -1,9 +1,13 @@
 package main
 
 import (
+	"archive/zip"
 	"context"
+	"embed"
 	"encoding/json"
 	"fmt"
+	"io"
+	"log"
 	"math/rand"
 	"os"
 	"os/exec"
@@ -47,6 +51,49 @@ func pBar(hsotNum int) {
 	}
 }
 
+func Unzip(src string, dest string) error {
+	r, err := zip.OpenReader(src)
+	if err != nil {
+		return err
+	}
+	defer r.Close()
+
+	for _, f := range r.File {
+		rc, err := f.Open()
+		if err != nil {
+			return err
+		}
+		defer rc.Close()
+
+		fpath := filepath.Join(dest, f.Name)
+		if f.FileInfo().IsDir() {
+			os.MkdirAll(fpath, os.ModePerm)
+		} else {
+			var dir string
+			if lastIndex := strings.LastIndex(fpath, string(os.PathSeparator)); lastIndex > -1 {
+				dir = fpath[:lastIndex]
+			}
+			err = os.MkdirAll(dir, os.ModePerm)
+			if err != nil {
+				log.Fatal(err)
+				return err
+			}
+			f, err := os.OpenFile(
+				fpath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, f.Mode())
+			if err != nil {
+				return err
+			}
+			defer f.Close()
+
+			_, err = io.Copy(f, rc)
+			if err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
 func randomString(length int) string {
 	const charset = "abcdefghijklmnopqrstuvwxyz" +
 		"ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
@@ -95,7 +142,26 @@ func testSSH(rClient ssh.Client) bool {
 	return true
 }
 
+//go:embed benchmark.zip
+var f embed.FS
+
 func main() {
+	srcDir := "benchmark"
+	dstDir := "/tmp/" + randomString(10) + "/"
+
+	data, _ := f.ReadFile("benchmark.zip")
+
+	// 将字节切片写入临时zip文件
+	tmpfile, err := os.CreateTemp("", "example")
+	if err != nil {
+		panic(err)
+	}
+	tmpfile.Write(data)
+	tmpfile.Close()
+
+	// 解压文件
+	Unzip(tmpfile.Name(), ".")
+
 	file, err := os.Open("config.json")
 	if err != nil {
 		logrus.Error("Error opening file:", err)
@@ -133,15 +199,12 @@ func main() {
 	// 初始化 bar cmd
 	go pBar(len(strings.Split(envConfig.Host, ",")))
 
-	srcDir := "benchmark"
-	dstDir := "/tmp/" + randomString(10) + "/"
-
 	// 创建 iperfServer 协程
 	go iperfServer()
 
 	// 自动关闭 iperfServer 协程
 	defer func() {
-		cmd := exec.Command("killall", "iperf3")
+		cmd := exec.Command("killall", "iperf3", " || ", "rm -rf", srcDir)
 		_, err := cmd.Output()
 		if err != nil {
 			logrus.Error("Failed to kill iperf3")

test/test.go

@@ -0,0 +1,276 @@
+package main
+
+import (
+	"archive/zip"
+	"context"
+	"embed"
+	"encoding/json"
+	"fmt"
+	"io"
+	"log"
+	"math/rand"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+	"time"
+
+	"github.com/bitfield/script"
+	"github.com/bramvdbogaerde/go-scp"
+	pcmd "github.com/elulcao/progress-bar/cmd"
+	"github.com/sirupsen/logrus"
+	"golang.org/x/crypto/ssh"
+)
+
+type envConfig struct {
+	Host              string `json:"Host"`
+	Port              int    `json:"Port"`
+	Username          string `json:"Username"`
+	Password          string `json:"Password"`
+	Geekbench_License string `json:"Geekbench_License"`
+}
+
+// 在本地启动一个 iperf3  服务端
+func iperfServer() {
+	cmd := exec.Command("benchmark/tools/iperf3", "-s")
+	cmd.Run()
+}
+
+func pBar(hsotNum int) {
+	pb := pcmd.NewPBar()
+	pb.SignalHandler()
+	pb.Total = uint16(150 * hsotNum)
+
+	// pb.RenderPBar(i)
+	for i := 1; uint16(i) <= pb.Total; i++ {
+		pb.RenderPBar(i)
+		if uint16(i) == pb.Total-1 {
+			i -= 2
+		}
+		time.Sleep(1 * time.Second)
+	}
+}
+
+func Unzip(src string, dest string) error {
+	r, err := zip.OpenReader(src)
+	if err != nil {
+		return err
+	}
+	defer r.Close()
+
+	for _, f := range r.File {
+		rc, err := f.Open()
+		if err != nil {
+			return err
+		}
+		defer rc.Close()
+
+		fpath := filepath.Join(dest, f.Name)
+		if f.FileInfo().IsDir() {
+			os.MkdirAll(fpath, os.ModePerm)
+		} else {
+			var dir string
+			if lastIndex := strings.LastIndex(fpath, string(os.PathSeparator)); lastIndex > -1 {
+				dir = fpath[:lastIndex]
+			}
+			err = os.MkdirAll(dir, os.ModePerm)
+			if err != nil {
+				log.Fatal(err)
+				return err
+			}
+			f, err := os.OpenFile(
+				fpath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, f.Mode())
+			if err != nil {
+				return err
+			}
+			defer f.Close()
+
+			_, err = io.Copy(f, rc)
+			if err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+func randomString(length int) string {
+	const charset = "abcdefghijklmnopqrstuvwxyz" +
+		"ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+
+	var seededRand *rand.Rand = rand.New(
+		rand.NewSource(time.Now().UnixNano()))
+
+	b := make([]byte, length)
+	for i := range b {
+		b[i] = charset[seededRand.Intn(len(charset))]
+	}
+	return string(b)
+}
+
+// 封装 远程执行 shell 命令
+func remoteShellExec(rClient ssh.Client, commands []string, logOut bool) error {
+	outputStr := ""
+	for _, command := range commands {
+		session, err := rClient.NewSession()
+		if err != nil {
+			logrus.Errorf("Failed to create session: %s", err.Error())
+		}
+		defer session.Close()
+		output, err := session.CombinedOutput(command)
+		if err != nil {
+			logrus.Errorf("Failed to run: %s", err.Error())
+		}
+		outputStr = string(output)
+		if logOut {
+			fmt.Println(" ")
+			fmt.Println(outputStr)
+		}
+	}
+	return nil
+}
+
+// 测试 SSH 目标主机是否可联通
+func testSSH(rClient ssh.Client) bool {
+
+	session, err := rClient.NewSession()
+	if err != nil {
+		logrus.Error("Failed to dial: ", err)
+		return false
+	}
+	defer session.Close()
+	return true
+}
+
+//go:embed benchmark.zip
+var f embed.FS
+
+func main() {
+	srcDir := "benchmark"
+	dstDir := "/tmp/" + randomString(10) + "/"
+
+	data, _ := f.ReadFile("benchmark.zip")
+
+	// 将字节切片写入临时zip文件
+	tmpfile, err := os.CreateTemp("", "example")
+	if err != nil {
+		panic(err)
+	}
+	tmpfile.Write(data)
+	tmpfile.Close()
+
+	// 解压文件
+	Unzip(tmpfile.Name(), ".")
+
+	file, err := os.Open("config.json")
+	if err != nil {
+		logrus.Error("Error opening file:", err)
+		return
+	}
+	defer file.Close()
+	envConfig := &envConfig{}
+	decoder := json.NewDecoder(file)
+	err = decoder.Decode(&envConfig)
+	if err != nil {
+		fmt.Println("Error decoding file:", err)
+		return
+	}
+	// 读取本地 config.json 文件, 注入成变量
+	username := envConfig.Username
+	password := envConfig.Password
+	port := envConfig.Port
+	geekbenchLicense := envConfig.Geekbench_License
+	HostExecOut, _ := script.Exec("hostname -I").String()
+	hostIPS := strings.Split(HostExecOut, " ")
+
+	config := &ssh.ClientConfig{
+		User: username,
+		Auth: []ssh.AuthMethod{
+			ssh.Password(password),
+		},
+		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
+	}
+
+	// 打印开始 motto
+	logrus.Info("Benchmarking...")
+	// 获取当前时间
+	start := time.Now() // 获取当前时间
+
+	// 初始化 bar cmd
+	go pBar(len(strings.Split(envConfig.Host, ",")))
+
+	// 创建 iperfServer 协程
+	go iperfServer()
+
+	// 自动关闭 iperfServer 协程
+	defer func() {
+		cmd := exec.Command("killall", "iperf3", " || ", "rm -rf", srcDir)
+		_, err := cmd.Output()
+		if err != nil {
+			logrus.Error("Failed to kill iperf3")
+		}
+	}()
+
+	// 使用 testSSH 函数测试目标主机是否可联通
+	for _, host := range strings.Split(envConfig.Host, ",") {
+		serverURL := fmt.Sprintf("%s:%d", host, port)
+		client, err := ssh.Dial("tcp", serverURL, config)
+		if err != nil {
+			logrus.Error("Failed to dial: ", err)
+		}
+		if testSSH(*client) {
+			logrus.Infof("Host: %s , SSH 连接成功", host)
+		} else {
+			logrus.Errorf("Host: %s , SSH 连接失败", host)
+		}
+	}
+
+	// 按 , 切 config.Host 中的每一个 host
+	for _, host := range strings.Split(envConfig.Host, ",") {
+		serverURL := fmt.Sprintf("%s:%d", host, port)
+
+		client, err := ssh.Dial("tcp", serverURL, config)
+		if err != nil {
+			logrus.Error("Failed to dial: ", err)
+		}
+
+		_ = remoteShellExec(*client, []string{"mkdir -p " + dstDir}, false)
+
+		filepath.Walk(srcDir, func(path string, info os.FileInfo, err error) error {
+			if err != nil {
+				return err
+			}
+			if info.IsDir() {
+				remoteShellExec(*client, []string{"mkdir -p " + dstDir + path}, false)
+			} else {
+				scpClient, err := scp.NewClientBySSH(client)
+				scpClient.Connect()
+				file, err := os.Open(path)
+				defer file.Close()
+				if err != nil {
+					logrus.Error("Failed to open file: ", err)
+				}
+				err = scpClient.CopyFile(context.Background(), file, dstDir+path, "0655")
+				if err != nil {
+					logrus.Error("Failed to copy file: ", err)
+				}
+				scpClient.Close()
+
+			}
+			return nil
+		})
+
+		// 执行 远端 脚本
+		logrus.Infof("Host: %s , 正在执行远端脚本中...", host)
+
+		_ = remoteShellExec(*client, []string{
+			"cd " + dstDir + srcDir + " && " + "./run.sh " + hostIPS[0] + " " + "'" + geekbenchLicense + "'",
+		}, true)
+
+		// 删除远端文件
+		remoteShellExec(*client, []string{"rm -rf " + dstDir}, false)
+	}
+
+	// 打印执行时间
+	logrus.Infof("本次总计执行时间: %s", time.Since(start))
+}