From 1fb0e79ea7a55f984d00989312e40bf50828b3ff Mon Sep 17 00:00:00 2001 From: "zun.yang" Date: Fri, 2 Jun 2023 10:48:30 +0800 Subject: [PATCH] Initial commit --- config.json | 6 ++ go.mod | 8 ++- go.sum | 32 ++++++++++ main.go | 180 +++++++++++++++++++++++++++++++++++++--------------- 4 files changed, 174 insertions(+), 52 deletions(-) create mode 100644 config.json diff --git a/config.json b/config.json new file mode 100644 index 0000000..fe21510 --- /dev/null +++ b/config.json @@ -0,0 +1,6 @@ +{ + "Host": "192.168.8.135,192.168.8.134", + "Port": 22, + "Username": "root", + "Password": "nuli123456" +} \ No newline at end of file diff --git a/go.mod b/go.mod index c2d8166..1b3b980 100644 --- a/go.mod +++ b/go.mod @@ -7,4 +7,10 @@ require ( golang.org/x/crypto v0.9.0 ) -require golang.org/x/sys v0.8.0 // indirect +require ( + github.com/bitfield/script v0.22.0 // indirect + github.com/itchyny/gojq v0.12.12 // indirect + github.com/itchyny/timefmt-go v0.1.5 // indirect + golang.org/x/sys v0.8.0 // indirect + mvdan.cc/sh/v3 v3.6.0 // indirect +) diff --git a/go.sum b/go.sum index 633479a..caebe9e 100644 --- a/go.sum +++ b/go.sum @@ -1,14 +1,46 @@ +github.com/bitfield/script v0.22.0 h1:LA7QHuEsXMPD52YLtxWrlqCCy+9FOpzNYfsRHC5Gsrc= +github.com/bitfield/script v0.22.0/go.mod h1:ms4w+9B8f2/W0mbsgWDVTtl7K94bYuZc3AunnJC4Ebs= github.com/bramvdbogaerde/go-scp v1.2.1 h1:BKTqrqXiQYovrDlfuVFaEGz0r4Ou6EED8L7jCXw6Buw= github.com/bramvdbogaerde/go-scp v1.2.1/go.mod h1:s4ZldBoRAOgUg8IrRP2Urmq5qqd2yPXQTPshACY8vQ0= +github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= +github.com/creack/pty v1.1.18/go.mod h1:MOBLtS5ELjhRRrroQr9kyvTxUAFNvYEK993ew/Vr4O4= +github.com/frankban/quicktest v1.14.4/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0= +github.com/google/go-cmp v0.5.4/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/google/renameio/v2 v2.0.0/go.mod h1:BtmJXm5YlszgC+TD4HOEEUFgkJP3nLxehU6hfe7jRt4= +github.com/itchyny/gojq v0.12.12 h1:x+xGI9BXqKoJQZkr95ibpe3cdrTbY8D9lonrK433rcA= +github.com/itchyny/gojq v0.12.12/go.mod h1:j+3sVkjxwd7A7Z5jrbKibgOLn0ZfLWkV+Awxr/pyzJE= +github.com/itchyny/timefmt-go v0.1.5 h1:G0INE2la8S6ru/ZI5JecgyzbbJNs5lG1RcBqa7Jm6GE= +github.com/itchyny/timefmt-go v0.1.5/go.mod h1:nEP7L+2YmAbT2kZ2HfSs1d8Xtw9LY8D2stDBckWakZ8= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/mattn/go-isatty v0.0.17/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM= +github.com/mattn/go-runewidth v0.0.14/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= +github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= +github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= +github.com/rivo/uniseg v0.4.4/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= +github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= +github.com/rogpeppe/go-internal v1.10.0/go.mod h1:UQnix2H7Ngw/k4C5ijL5+65zddjncjaFoBhdsK/akog= golang.org/x/crypto v0.0.0-20210513164829-c07d793c2f9a/go.mod h1:P+XmwS30IXTQdn5tA2iutPOUgjI07+tq3H3K9MVA1s8= golang.org/x/crypto v0.9.0 h1:LF6fAI+IutBocDJ2OT0Q1g8plpYljMZ4+lty+dsqw3g= golang.org/x/crypto v0.9.0/go.mod h1:yrmDGqONDYtNj3tH8X9dzUun2m2lzPa9ngI6/RUPGR0= golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg= +golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210525143221-35b2ab0089ea/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.3.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.8.0 h1:EBmGv8NaZBZTWvrbjNoL6HVt+IVy3QDQpJs7VRIw3tU= golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= +golang.org/x/term v0.3.0/go.mod h1:q750SLmJuPmVoN1blW3UFBPREJfb1KmY3vwxfr+nFDA= golang.org/x/term v0.8.0 h1:n5xxQn2i3PC0yLAbjTpNT85q/Kgzcr2gIoX9OrJUols= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= +golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +mvdan.cc/editorconfig v0.2.0/go.mod h1:lvnnD3BNdBYkhq+B4uBuFFKatfp02eB6HixDvEz91C0= +mvdan.cc/sh/v3 v3.6.0 h1:gtva4EXJ0dFNvl5bHjcUEvws+KRcDslT8VKheTYkbGU= +mvdan.cc/sh/v3 v3.6.0/go.mod h1:U4mhtBLZ32iWhif5/lD+ygy1zrgaQhUu+XFy7C8+TTA= diff --git a/main.go b/main.go index 1a58fee..3d2c5e2 100644 --- a/main.go +++ b/main.go @@ -2,28 +2,107 @@ package main import ( "context" + "encoding/json" "fmt" "log" + "math/rand" "os" "os/exec" "path/filepath" + "strings" + "time" - scp "github.com/bramvdbogaerde/go-scp" + "github.com/bitfield/script" + "github.com/bramvdbogaerde/go-scp" "golang.org/x/crypto/ssh" ) +type envConfig struct { + Host string `json:"Host"` + Port int `json:"Port"` + Username string `json:"Username"` + Password string `json:"Password"` +} + +// 在本地启动一个 iperf3 服务端 func iperfServer() { - // 在本地启动一个 iperf3 服务端 cmd := exec.Command("benchmark/tools/iperf3", "-s") cmd.Run() } +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 { + log.Fatal("Failed to create session: ", err) + } + defer session.Close() + output, err := session.CombinedOutput(command) + if err != nil { + log.Fatal("Failed to run: ", err) + } + outputStr = string(output) + if logOut { + log.Println(outputStr) + } + } + return nil +} + +// 测试 SSH 目标主机是否可联通 +func testSSH(rClient ssh.Client) bool { + + session, err := rClient.NewSession() + if err != nil { + log.Fatal("Failed to dial: ", err) + return false + } + defer session.Close() + return true +} + func main() { - srcDir := "./benchmark" - dstDir := "/tmp/benchmark" - username := "root" - password := "nuli123456" - serverURL := "192.168.8.135:22" + file, err := os.Open("config.json") + if err != nil { + fmt.Println("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 + } + + srcDir := "benchmark" + dstDir := "/tmp/" + randomString(10) + "/" + + // 读取本地 config.json 文件, 注入成变量 + + username := envConfig.Username + password := envConfig.Password + port := envConfig.Port + HostExecOut, _ := script.Exec("hostname -I").String() + hostIPS := strings.Split(HostExecOut, " ") config := &ssh.ClientConfig{ User: username, @@ -45,64 +124,63 @@ func main() { } }() - client, err := ssh.Dial("tcp", serverURL, config) - if err != nil { - log.Fatal("Failed to dial: ", err) - } - - commands := []string{ - "mkdir -p " + dstDir, - } - - for _, command := range commands { - session, err := client.NewSession() + // 使用 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 { - log.Fatal("Failed to create session: ", err) + log.Fatal("Failed to dial: ", err) } - defer session.Close() - - _, err = session.CombinedOutput(command) - if err != nil { - log.Fatal("Failed to run: ", err) + if testSSH(*client) { + log.Printf("Host: %s , SSH 连接成功", host) + } else { + log.Printf("Host: %s , SSH 连接失败", host) } } - // Create a new SCP client - scpClient, err := scp.NewClientBySSH(client) + // 按 , 切 config.Host 中的每一个 host + for _, host := range strings.Split(envConfig.Host, ",") { + serverURL := fmt.Sprintf("%s:%d", host, port) - if err != nil { - fmt.Println("Error creating new SSH session from existing connection", err) - } - - // Walk the source directory - err = filepath.Walk(srcDir, func(path string, info os.FileInfo, err error) error { + client, err := ssh.Dial("tcp", serverURL, config) if err != nil { - return err + log.Fatal("Failed to dial: ", err) } - if !info.IsDir() { - // Open the source file - f, err := os.Open(path) + _ = remoteShellExec(*client, []string{"mkdir -p " + dstDir}, false) + + filepath.Walk(srcDir, func(path string, info os.FileInfo, err error) error { if err != nil { return err } - defer f.Close() + 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 { + log.Fatal("Failed to open file: ", err) + } + err = scpClient.CopyFile(context.Background(), file, dstDir+path, "0655") + if err != nil { + log.Fatal("Failed to copy file: ", err) + } + scpClient.Close() - // Calculate the destination path - relPath, _ := filepath.Rel(srcDir, path) - dstPath := filepath.Join(dstDir, relPath) - - // Copy the file to the remote server - err = scpClient.CopyFile(context.Background(), f, dstPath, "0655") - if err != nil { - return err } - } - return nil - }) + return nil + }) - if err != nil { - log.Fatalf("Failed to copy directory: %v", err) + // 执行 远端 脚本 + log.Printf("Host: %s , 正在执行远端脚本...", host) + + _ = remoteShellExec(*client, []string{ + "cd " + dstDir + srcDir + " && " + "./run.sh " + hostIPS[0], + }, true) + + // 删除远端文件 + remoteShellExec(*client, []string{"rm -rf " + dstDir}, false) } - }