Go 连接服务器实现交互，并将命令执行结果写入文件

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) {
	// 创建ssh登录配置
	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
	}

	// dial连接服务器
	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) // 秘钥登录(本机密钥需未设置密码，否则无法登陆)
	//conn, err := keyOrPwdConnect("***.***.***.***", "root", "123456", "", 22)                  // 密码登录  
	if err != nil {
		panic(err)
	}
	defer conn.Close()

	// 创建 ssh session 会话
	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)
		}		
	}
}