be.ems/src/framework/utils/ssh/ssh.go

package ssh

import (
	"bytes"
	"fmt"
	"io"
	"strings"
	"sync"
	"time"

	gossh "golang.org/x/crypto/ssh"
)

// ConnSSH 连接SSH对象
type ConnSSH struct {
	User       string `json:"user"`       // 主机用户名
	Addr       string `json:"addr"`       // 主机地址
	Port       int    `json:"port"`       // SSH端口
	AuthMode   string `json:"authMode"`   // 认证模式（0密码 1主机私钥）
	Password   string `json:"password"`   // 认证密码
	PrivateKey string `json:"privateKey"` // 认证私钥
	PassPhrase string `json:"passPhrase"` // 认证私钥密码

	DialTimeOut time.Duration `json:"dialTimeOut"` // 连接超时断开

	Client     *gossh.Client `json:"client"`
	LastResult string        `json:"lastResult"` // 记最后一次执行命令的结果
}

// NewClient 创建SSH客户端
func (c *ConnSSH) NewClient() (*ConnSSH, error) {
	// IPV6地址协议
	proto := "tcp"
	if strings.Contains(c.Addr, ":") {
		proto = "tcp6"
		c.Addr = fmt.Sprintf("[%s]", c.Addr)
	}
	addr := fmt.Sprintf("%s:%d", c.Addr, c.Port)

	// ssh客户端配置
	config := &gossh.ClientConfig{}
	config.SetDefaults()
	config.HostKeyCallback = gossh.InsecureIgnoreHostKey()
	config.User = c.User

	// 默认等待5s
	if c.DialTimeOut == 0 {
		c.DialTimeOut = 5 * time.Second
	}
	config.Timeout = c.DialTimeOut

	// 认证模式-0密码 1私钥
	if c.AuthMode == "1" {
		var signer gossh.Signer
		var err error
		if len(c.PassPhrase) != 0 {
			signer, err = gossh.ParsePrivateKeyWithPassphrase([]byte(c.PrivateKey), []byte(c.PassPhrase))
		} else {
			signer, err = gossh.ParsePrivateKey([]byte(c.PrivateKey))
		}
		if err != nil {
			return nil, err
		}
		config.Auth = []gossh.AuthMethod{gossh.PublicKeys(signer)}
	} else {
		config.Auth = []gossh.AuthMethod{gossh.Password(c.Password)}
	}

	client, err := gossh.Dial(proto, addr, config)
	if nil != err {
		return c, err
	}
	c.Client = client
	return c, nil
}

// Close 关闭当前SSH客户端
func (c *ConnSSH) Close() {
	if c.Client != nil {
		c.Client.Close()
	}
}

// RunCMD 执行单次命令
func (c *ConnSSH) RunCMD(cmd string) (string, error) {
	if c.Client == nil {
		if _, err := c.NewClient(); err != nil {
			return "", err
		}
	}
	session, err := c.Client.NewSession()
	if err != nil {
		return "", err
	}
	defer session.Close()
	buf, err := session.CombinedOutput(cmd)

	c.LastResult = string(buf)
	return c.LastResult, err
}

// NewClient 创建SSH客户端会话对象
func (c *ConnSSH) NewClientSession(cols, rows int) (*SSHClientSession, error) {
	sshSession, err := c.Client.NewSession()
	if err != nil {
		return nil, err
	}

	stdin, err := sshSession.StdinPipe()
	if err != nil {
		return nil, err
	}

	comboWriter := new(singleWriter)
	sshSession.Stdout = comboWriter
	sshSession.Stderr = comboWriter

	modes := gossh.TerminalModes{
		gossh.ECHO:          1,
		gossh.TTY_OP_ISPEED: 14400,
		gossh.TTY_OP_OSPEED: 14400,
	}

	if err := sshSession.RequestPty("xterm", rows, cols, modes); err != nil {
		return nil, err
	}

	if err := sshSession.Shell(); err != nil {
		return nil, err
	}

	return &SSHClientSession{
		Stdin:   stdin,
		Stdout:  comboWriter,
		Session: sshSession,
	}, nil
}

// SSHClientSession SSH客户端会话对象
type SSHClientSession struct {
	Stdin   io.WriteCloser
	Stdout  *singleWriter
	Session *gossh.Session
}

// Close 关闭会话
func (s *SSHClientSession) Close() {
	if s.Stdin != nil {
		s.Stdin.Close()
	}
	if s.Stdout != nil {
		s.Stdout = nil
	}
	if s.Session != nil {
		s.Session.Close()
	}
}

// Write 写入命令 回车(\n)才会执行
func (s *SSHClientSession) Write(cmd string) (int, error) {
	if s.Stdin == nil {
		return 0, fmt.Errorf("ssh client session is nil to content write failed")
	}
	return s.Stdin.Write([]byte(cmd))
}

// Read 读取结果 等待一会才有结果
func (s *SSHClientSession) Read() []byte {
	if s.Stdout == nil {
		return []byte{}
	}
	// time.Sleep(300 * time.Millisecond)
	bs := s.Stdout.Bytes()
	if len(bs) > 0 {
		s.Stdout.Reset()
		return bs
	}
	return []byte{}
}

// CombinedOutput 发送命令带结果返回
func (s *SSHClientSession) CombinedOutput(cmd string) (string, error) {
	n, err := s.Write(cmd)
	if n == 0 || err != nil {
		return "", err
	}
	return string(s.Read()), nil
}

// singleWriter SSH客户端会话消息
type singleWriter struct {
	b  bytes.Buffer
	mu sync.Mutex
}

func (w *singleWriter) Write(p []byte) (int, error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	return w.b.Write(p)
}
func (w *singleWriter) Bytes() []byte {
	w.mu.Lock()
	defer w.mu.Unlock()
	return w.b.Bytes()
}
func (w *singleWriter) Reset() {
	w.mu.Lock()
	defer w.mu.Unlock()
	w.b.Reset()
}