OMC/be.ems
1
0
Fork 0
Code Issues Releases 3 Activity

style: 调整ssh工具包代码分割

Browse Source
This commit is contained in:
TsMask
2024-06-14 16:55:27 +08:00
parent e254f0b1f9
commit 3fa72ae983
3 changed files with 326 additions and 310 deletions
Download Patch File Download Diff File Expand all files Collapse all files

361
src/framework/utils/ssh/ssh.go
View File

@@ -1,14 +1,10 @@
package ssh
import (
"bytes"
"fmt"
"io"
"os"
"os/user"
"path/filepath"
"strings"
"sync"
"time"
"be.ems/src/framework/logger"
@@ -89,24 +85,6 @@ func (c *ConnSSH) Close() {
}
}
// NewClientByLocalPrivate 创建SSH客户端-本地私钥(~/.ssh/id_rsa)直连
//
// ssh.ConnSSH{
// User: "user",
// Addr: "192.168.x.x",
// Port: body.Port,
// }
func (c *ConnSSH) NewClientByLocalPrivate() (*ConnSSH, error) {
c.Port = 22
c.AuthMode = "1"
privateKey, err := c.CurrentUserRsaKey(false)
if err != nil {
return nil, err
}
c.PrivateKey = privateKey
return c.NewClient()
}
// RunCMD 执行单次命令
func (c *ConnSSH) RunCMD(cmd string) (string, error) {
if c.Client == nil {
@@ -128,57 +106,6 @@ func (c *ConnSSH) RunCMD(cmd string) (string, error) {
return c.LastResult, err
}
// SendToAuthorizedKeys 发送当前用户私钥到远程服务器进行授权密钥
func (c *ConnSSH) SendToAuthorizedKeys() error {
publicKey, err := c.CurrentUserRsaKey(true)
if err != nil {
return err
}
authorizedKeysEntry := fmt.Sprintln(strings.TrimSpace(publicKey))
cmdStrArr := []string{
fmt.Sprintf("sudo mkdir -p /home/%s/.ssh && sudo chown %s:%s /home/%s/.ssh && sudo chmod 700 /home/%s/.ssh", c.User, c.User, c.User, c.User, c.User),
fmt.Sprintf("sudo touch /home/%s/.ssh/authorized_keys && sudo chown %s:%s /home/%s/.ssh/authorized_keys && sudo chmod 600 /home/%s/.ssh/authorized_keys", c.User, c.User, c.User, c.User, c.User),
fmt.Sprintf("echo '%s' | sudo tee -a /home/%s/.ssh/authorized_keys", authorizedKeysEntry, c.User),
}
_, err = c.RunCMD(strings.Join(cmdStrArr, " && "))
if err != nil {
logger.Errorf("SendAuthorizedKeys echo err %s", err.Error())
return err
}
return nil
}
// 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
}
// NewClientSession 创建SSH客户端会话对象
func (c *ConnSSH) NewClientSession(cols, rows int) (*SSHClientSession, error) {
sshSession, err := c.Client.NewSession()
@@ -216,69 +143,6 @@ func (c *ConnSSH) NewClientSession(cols, rows int) (*SSHClientSession, error) {
}, 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{}
}
bs := s.Stdout.Bytes()
if len(bs) > 0 {
s.Stdout.Reset()
return bs
}
return []byte{}
}
// 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()
}
// NewClientSFTP 创建SSH客户端SFTP对象
func (c *ConnSSH) NewClientSFTP() (*SSHClientSFTP, error) {
sftpClient, err := gosftp.NewClient(c.Client)
@@ -292,193 +156,70 @@ func (c *ConnSSH) NewClientSFTP() (*SSHClientSFTP, error) {
}, nil
}
// SSHClientSFTP SSH客户端SFTP对象
type SSHClientSFTP struct {
Client *gosftp.Client
// NewClientByLocalPrivate 创建SSH客户端-本地私钥(~/.ssh/id_rsa)直连
//
// ssh.ConnSSH{
// User: "user",
// Addr: "192.168.x.x",
// Port: body.Port,
// }
func (c *ConnSSH) NewClientByLocalPrivate() (*ConnSSH, error) {
c.Port = 22
c.AuthMode = "1"
privateKey, err := c.CurrentUserRsaKey(false)
if err != nil {
return nil, err
}
c.PrivateKey = privateKey
return c.NewClient()
}
// Close 关闭会话
func (s *SSHClientSFTP) Close() {
if s.Client != nil {
s.Client.Close()
}
}
// CopyDirRemoteToLocal 复制目录-远程到本地
func (s *SSHClientSFTP) CopyDirRemoteToLocal(remoteDir, localDir string) error {
// 列出远程目录中的文件和子目录
remoteFiles, err := s.Client.ReadDir(remoteDir)
// 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("CopyDirRemoteToLocal failed to reading remote directory %s: => %s", remoteDir, err.Error())
return err
logger.Errorf("CurrentUserRsaKey get => %s", err.Error())
return "", err
}
// 创建本地目录
err = os.MkdirAll(localDir, 0775)
if err != nil {
logger.Errorf("CopyDirRemoteToLocal failed to creating local directory %s: => %s", localDir, err.Error())
return err
}
// 遍历远程文件和子目录并复制到本地
for _, remoteFile := range remoteFiles {
remotePath := filepath.Join(remoteDir, remoteFile.Name())
localPath := filepath.Join(localDir, remoteFile.Name())
if remoteFile.IsDir() {
// 如果是子目录，则递归复制子目录
err = s.CopyDirRemoteToLocal(remotePath, localPath)
if err != nil {
logger.Errorf("CopyDirRemoteToLocal failed to copying remote directory %s: => %s", remotePath, err.Error())
continue
}
} else {
// 如果是文件，则复制文件内容
remoteFile, err := s.Client.Open(remotePath)
if err != nil {
logger.Errorf("CopyDirRemoteToLocal failed to opening remote file %s: => %s", remotePath, err.Error())
continue
}
defer remoteFile.Close()
localFile, err := os.Create(localPath)
if err != nil {
logger.Errorf("CopyDirRemoteToLocal failed to creating local file %s: => %s", localPath, err.Error())
continue
}
defer localFile.Close()
_, err = io.Copy(localFile, remoteFile)
if err != nil {
logger.Errorf("CopyDirRemoteToLocal failed to copying file contents from %s to %s: => %s", remotePath, localPath, err.Error())
continue
}
// 是否存在私钥并创建
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())
}
}
return nil
// 读取用户默认的文件
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
}
// CopyDirRemoteToLocal 复制目录-本地到远程
func (s *SSHClientSFTP) CopyDirLocalToRemote(localDir, remoteDir string) error {
// 创建远程目录
err := s.Client.MkdirAll(remoteDir)
// SendToAuthorizedKeys 发送当前用户私钥到远程服务器进行授权密钥
func (c *ConnSSH) SendToAuthorizedKeys() error {
publicKey, err := c.CurrentUserRsaKey(true)
if err != nil {
logger.Errorf("CopyDirLocalToRemote failed to creating remote directory %s: => %s", remoteDir, err.Error())
return err
}
// 遍历本地目录中的文件和子目录并复制到远程
err = filepath.Walk(localDir, func(localPath string, info os.FileInfo, err error) error {
if err != nil {
return err
}
// 生成远程路径
remotePath := filepath.Join(remoteDir, localPath[len(localDir):])
if info.IsDir() {
// 如果是子目录，则创建远程目录
err := s.Client.MkdirAll(remotePath)
if err != nil {
logger.Errorf("CopyDirLocalToRemote failed to creating remote directory %s: => %s", remotePath, err.Error())
return nil
}
} else {
// 如果是文件，则复制文件内容
localFile, err := os.Open(localPath)
if err != nil {
logger.Errorf("CopyDirLocalToRemote failed to opening local file %s: => %s", localPath, err.Error())
return nil
}
defer localFile.Close()
remoteFile, err := s.Client.Create(remotePath)
if err != nil {
logger.Errorf("CopyDirLocalToRemote failed to creating remote file %s: => %s", remotePath, err.Error())
return nil
}
defer remoteFile.Close()
_, err = io.Copy(remoteFile, localFile)
if err != nil {
logger.Errorf("CopyDirLocalToRemote failed to copying file contents from %s to %s: => %s", localPath, remotePath, err.Error())
return nil
}
}
return nil
})
authorizedKeysEntry := fmt.Sprintln(strings.TrimSpace(publicKey))
cmdStrArr := []string{
fmt.Sprintf("sudo mkdir -p /home/%s/.ssh && sudo chown %s:%s /home/%s/.ssh && sudo chmod 700 /home/%s/.ssh", c.User, c.User, c.User, c.User, c.User),
fmt.Sprintf("sudo touch /home/%s/.ssh/authorized_keys && sudo chown %s:%s /home/%s/.ssh/authorized_keys && sudo chmod 600 /home/%s/.ssh/authorized_keys", c.User, c.User, c.User, c.User, c.User),
fmt.Sprintf("echo '%s' | sudo tee -a /home/%s/.ssh/authorized_keys", authorizedKeysEntry, c.User),
}
_, err = c.RunCMD(strings.Join(cmdStrArr, " && "))
if err != nil {
logger.Errorf("CopyDirLocalToRemote failed to walking local directory: => %s", err.Error())
return err
}
return nil
}
// CopyDirRemoteToLocal 复制文件-远程到本地
func (s *SSHClientSFTP) CopyFileRemoteToLocal(remotePath, localPath string) error {
// 打开远程文件
remoteFile, err := s.Client.Open(remotePath)
if err != nil {
logger.Errorf("CopyFileRemoteToLocal failed to opening remote file: => %s", err.Error())
return err
}
defer remoteFile.Close()
if err := os.MkdirAll(filepath.Dir(localPath), 0775); err != nil {
return err
}
// 如果目标文件已经存在，先将目标文件重命名
// if info, err := os.Stat(localPath); err == nil && !info.IsDir() {
// ext := filepath.Ext(localPath)
// name := localPath[0 : len(localPath)-len(ext)]
// newName := fmt.Sprintf("%s-%s%s", name, time.Now().Format("20060102_150405"), ext)
// err := os.Rename(localPath, newName)
// if err != nil {
// return err
// }
// }
// 创建本地文件
localFile, err := os.Create(localPath)
if err != nil {
logger.Errorf("CopyFileRemoteToLocal failed to creating local file: => %s", err.Error())
return err
}
defer localFile.Close()
// 复制文件内容
_, err = io.Copy(localFile, remoteFile)
if err != nil {
logger.Errorf("CopyFileRemoteToLocal failed to copying contents: => %s", err.Error())
return err
}
return nil
}
// CopyDirRemoteToLocal 复制文件-本地到远程
func (s *SSHClientSFTP) CopyFileLocalToRemote(localPath, remotePath string) error {
// 打开本地文件
localFile, err := os.Open(localPath)
if err != nil {
logger.Errorf("CopyFileLocalToRemote failed to opening local file: => %s", err.Error())
return err
}
defer localFile.Close()
// 创建远程文件
remoteFile, err := s.Client.Create(remotePath)
if err != nil {
logger.Errorf("CopyFileLocalToRemote failed to creating remote file: => %s", err.Error())
return err
}
defer remoteFile.Close()
// 复制文件内容
_, err = io.Copy(remoteFile, localFile)
if err != nil {
logger.Errorf("CopyFileLocalToRemote failed to copying contents: => %s", err.Error())
logger.Errorf("SendAuthorizedKeys echo err %s", err.Error())
return err
}
return nil