From 455e980af62e44f583ac2f32c0e654ee036a2fbd Mon Sep 17 00:00:00 2001 From: "zun.yang" Date: Mon, 31 Jul 2023 16:47:55 +0800 Subject: [PATCH] Update --- .gitignore | 2 +- main.go | 71 ++++++++++++- test/test.go | 276 +++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 344 insertions(+), 5 deletions(-) create mode 100644 test/test.go diff --git a/.gitignore b/.gitignore index 20d9c71..290184e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,7 +1,7 @@ bin/ vendor ssh-benchmark - +benchmark.zip --- .vscode/* !.vscode/settings.json diff --git a/main.go b/main.go index db72b30..23bae36 100644 --- a/main.go +++ b/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") diff --git a/test/test.go b/test/test.go new file mode 100644 index 0000000..23bae36 --- /dev/null +++ b/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)) +}