Go 连接服务器实现交互,并将命令执行结果写入文件
做性能测试,需要记录相关的指标,于是弄了一个自动化脚本,可以实现
相关指标存入文件,方便测试人员记录测试数据;
有部分代码是对输出结果字符串进行处理,可根据自己实际需求进行处理;
说明:本脚本部分代码借鉴与其他开发者,如有侵权,请联系作者删帖;
package main
import (
"errors"
"fmt"
"log"
"os"
"strings"
"regexp"
"io/ioutil"
"golang.org/x/crypto/ssh"
)
func keyOrPwdConnect(sshHost, sshUser, sshPassword, sshKey string, sshPort int) (*ssh.Client, error) {
config := &ssh.ClientConfig{
User: sshUser,
HostKeyCallback: ssh.InsecureIgnoreHostKey(),
}
if sshKey != "" {
key, err := ioutil.ReadFile(sshKey)
if err != nil {
panic("秘钥读取失败")
}
signer, err := ssh.ParsePrivateKey(key)
if err != nil {
panic("秘钥签名失败")
}
config.Auth = []ssh.AuthMethod{
ssh.PublicKeys(signer),
}
} else if sshPassword != "" {
config.Auth = []ssh.AuthMethod{ssh.Password(sshPassword)}
} else {
err := errors.New("秘钥或者密码登录")
return nil, err
}
addr := fmt.Sprintf("%s:%d", sshHost, sshPort)
Client, err := ssh.Dial("tcp", addr, config)
if err != nil {
log.Fatal("连接到服务器失败", err)
return nil, err
}
return Client, nil
}
func delete_extra_space(count string) string {
s1 := strings.Replace(count, " ", " ", -1)
regstr := "\\s{2,}"
reg, _ := regexp.Compile(regstr)
s2 := make([]byte, len(s1))
copy(s2, s1)
spc_index := reg.FindStringIndex(string(s2))
for len(spc_index) > 0 {
s2 = append(s2[:spc_index[0]+1], s2[spc_index[1]:]...)
spc_index = reg.FindStringIndex(string(s2))
}
return string(s2)
}
func writeToFile(str string) {
path := "test/test.txt"
f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
if err != nil {
log.Println(err.Error())
}
_, err = f.Write([]byte(str))
if err != nil {
log.Println(err.Error())
}
f.Close()
}
func main() {
conn, err := keyOrPwdConnect("***.***.***.***", "root", "", "~/.ssh/id_rsa", 22)
if err != nil {
panic(err)
}
defer conn.Close()
session, err := conn.NewSession()
if err != nil {
panic(err)
}
defer session.Close()
cmd := "watch -n 1 nvidia-smi"
cmdInfo, err := session.CombinedOutput(cmd)
if err != nil {
panic(err)
}
count := string (cmdInfo)
lines := strings.Split(count, "\n")
for index, line := range lines {
if index == 0 {
continue
}
words := strings.Split(delete_extra_space(line), " ")
if len(words) >2 {
str := words[1] + " " + words[2]+"\n"
writeToFile(str)
}
}
}