SSH-Benchmark/main.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"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 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
}

func main() {
	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, ",")))

	srcDir := "benchmark"
	dstDir := "/tmp/" + randomString(10) + "/"

	// 创建 iperfServer 协程
	go iperfServer()

	// 自动关闭 iperfServer 协程
	defer func() {
		cmd := exec.Command("killall", "iperf3")
		_, 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))
}