SSH-Benchmark/main.go
2023-06-01 14:05:07 +08:00

109 lines
2.0 KiB
Go

package main
import (
"context"
"fmt"
"log"
"os"
"os/exec"
"path/filepath"
scp "github.com/bramvdbogaerde/go-scp"
"golang.org/x/crypto/ssh"
)
func iperfServer() {
// 在本地启动一个 iperf3 服务端
cmd := exec.Command("benchmark/tools/iperf3", "-s")
cmd.Run()
}
func main() {
srcDir := "./benchmark"
dstDir := "/tmp/benchmark"
username := "root"
password := "nuli123456"
serverURL := "192.168.8.135:22"
config := &ssh.ClientConfig{
User: username,
Auth: []ssh.AuthMethod{
ssh.Password(password),
},
HostKeyCallback: ssh.InsecureIgnoreHostKey(),
}
// 创建 iperfServer 协程
go iperfServer()
// 自动关闭 iperfServer 协程
defer func() {
cmd := exec.Command("killall", "iperf3")
_, err := cmd.Output()
if err != nil {
log.Fatal(err)
}
}()
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()
if err != nil {
log.Fatal("Failed to create session: ", err)
}
defer session.Close()
_, err = session.CombinedOutput(command)
if err != nil {
log.Fatal("Failed to run: ", err)
}
}
// Create a new SCP client
scpClient, err := scp.NewClientBySSH(client)
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 {
if err != nil {
return err
}
if !info.IsDir() {
// Open the source file
f, err := os.Open(path)
if err != nil {
return err
}
defer f.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
})
if err != nil {
log.Fatalf("Failed to copy directory: %v", err)
}
}