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

package ssh

import (
	"fmt"
	"os"
	"os/user"
	"strings"
	"time"

	"be.ems/src/framework/logger"
	"be.ems/src/framework/utils/cmd"
	gosftp "github.com/pkg/sftp"
	gossh "golang.org/x/crypto/ssh"
)

// ConnSSH 连接SSH对象
type ConnSSH struct {
	User       string `json:"user"`       // 主机用户名
	Addr       string `json:"addr"`       // 主机地址
	Port       int64  `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 {
			logger.Errorf("NewClient parse private key => %s", err.Error())
			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 {
		logger.Errorf("NewClient dial => %s", err.Error())
		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 {
		logger.Errorf("RunCMD failed to create session: => %s", err.Error())
		return "", err
	}
	defer session.Close()
	buf, err := session.CombinedOutput(cmd)
	if err != nil {
		logger.Infof("RunCMD failed run command: => %s", cmd)
		logger.Errorf("RunCMD failed run command: => %s", err.Error())
	}
	c.LastResult = string(buf)
	return c.LastResult, err
}

// NewClientSession 创建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
}

// NewClientSFTP 创建SSH客户端SFTP对象
func (c *ConnSSH) NewClientSFTP() (*SSHClientSFTP, error) {
	sftpClient, err := gosftp.NewClient(c.Client)
	if err != nil {
		logger.Errorf("NewClientSFTP failed to create sftp: => %s", err.Error())
		return nil, err
	}

	return &SSHClientSFTP{
		Client: sftpClient,
	}, nil
}

// NewClientByLocalPrivate 创建SSH客户端-本地私钥(~/.ssh/id_rsa)直连
//
//	ssh.ConnSSH{
//		User: "user",
//		Addr: "192.168.x.x",
//		Port: body.Port,
//	}
func (c *ConnSSH) NewClientByLocalPrivate() (*ConnSSH, error) {
	c.AuthMode = "1"
	privateKey, err := c.CurrentUserRsaKey(false)
	if err != nil {
		return nil, err
	}
	c.PrivateKey = privateKey
	return c.NewClient()
}

// CurrentUserRsaKey 当前用户OMC使用的RSA私钥
// 默认读取私钥
// ssh-keygen -t rsa -P "" -f ~/.ssh/id_rsa
// ssh-keygen -y -f ~/.ssh/id_rsa > ~/.ssh/id_rsa.pub
func (c *ConnSSH) CurrentUserRsaKey(publicKey bool) (string, error) {
	usr, err := user.Current()
	if err != nil {
		logger.Errorf("CurrentUserRsaKey get => %s", err.Error())
		return "", err
	}

	// 是否存在私钥并创建
	keyPath := fmt.Sprintf("%s/.ssh/id_rsa", usr.HomeDir)
	if _, err := os.Stat(keyPath); err != nil {
		if _, err := cmd.ExecWithCheck("ssh-keygen", "-t", "rsa", "-P", "", "-f", keyPath); err != nil {
			logger.Errorf("CurrentUserPrivateKey ssh-keygen [%s] rsa => %s", usr.Username, err.Error())
		}
	}

	// 读取用户默认的文件
	if publicKey {
		keyPath = keyPath + ".pub"
	}
	key, err := os.ReadFile(keyPath)
	if err != nil {
		logger.Errorf("CurrentUserRsaKey [%s] read => %s", usr.Username, err.Error())
		return "", fmt.Errorf("read file %s fail", keyPath)
	}
	return string(key), nil
}

// SendToAuthorizedKeys 发送当前用户私钥到远程服务器进行授权密钥
func (c *ConnSSH) SendToAuthorizedKeys() error {
	publicKey, err := c.CurrentUserRsaKey(true)
	if err != nil {
		return err
	}
	// "sudo mkdir -p ~/.ssh && sudo chown omcuser:omcuser ~/.ssh && sudo chmod 700 ~/.ssh"
	// "sudo touch ~/.ssh/authorized_keys && sudo chown omcuser:omcuser ~/.ssh/authorized_keys && sudo chmod 600 ~/.ssh/authorized_keys"
	// "echo 'ssh-rsa AAAAB3= pc-host\n' | sudo tee -a ~/.ssh/authorized_keys"
	authorizedKeysEntry := fmt.Sprintln(strings.TrimSpace(publicKey))
	cmdStrArr := []string{
		fmt.Sprintf("sudo mkdir -p ~/.ssh && sudo chown %s:%s ~/.ssh && sudo chmod 700 ~/.ssh", c.User, c.User),
		fmt.Sprintf("sudo touch ~/.ssh/authorized_keys && sudo chown %s:%s ~/.ssh/authorized_keys && sudo chmod 600 ~/.ssh/authorized_keys", c.User, c.User),
		fmt.Sprintf("echo '%s' | sudo tee -a ~/.ssh/authorized_keys", authorizedKeysEntry),
	}
	_, err = c.RunCMD(strings.Join(cmdStrArr, " && "))
	if err != nil {
		logger.Errorf("SendAuthorizedKeys echo err %s", err.Error())
		return err
	}
	return nil
}