From 3fa72ae9837443c222cf0f9ebc4f212bb4ac020e Mon Sep 17 00:00:00 2001 From: TsMask <340112800@qq.com> Date: Fri, 14 Jun 2024 16:55:27 +0800 Subject: [PATCH] =?UTF-8?q?style:=20=E8=B0=83=E6=95=B4ssh=E5=B7=A5?= =?UTF-8?q?=E5=85=B7=E5=8C=85=E4=BB=A3=E7=A0=81=E5=88=86=E5=89=B2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/framework/utils/ssh/sftp.go | 202 ++++++++++++++ src/framework/utils/ssh/ssh.go | 361 ++++--------------------- src/framework/utils/ssh/ssh_session.go | 73 +++++ 3 files changed, 326 insertions(+), 310 deletions(-) create mode 100644 src/framework/utils/ssh/sftp.go create mode 100644 src/framework/utils/ssh/ssh_session.go diff --git a/src/framework/utils/ssh/sftp.go b/src/framework/utils/ssh/sftp.go new file mode 100644 index 00000000..33367337 --- /dev/null +++ b/src/framework/utils/ssh/sftp.go @@ -0,0 +1,202 @@ +package ssh + +import ( + "io" + "os" + "path/filepath" + + "be.ems/src/framework/logger" + gosftp "github.com/pkg/sftp" +) + +// SSHClientSFTP SSH客户端SFTP对象 +type SSHClientSFTP struct { + Client *gosftp.Client +} + +// 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) + if err != nil { + logger.Errorf("CopyDirRemoteToLocal failed to reading remote directory %s: => %s", remoteDir, 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 + } + } + } + return nil +} + +// CopyDirRemoteToLocal 复制目录-本地到远程 +func (s *SSHClientSFTP) CopyDirLocalToRemote(localDir, remoteDir string) error { + // 创建远程目录 + err := s.Client.MkdirAll(remoteDir) + 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 + }) + 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()) + return err + } + return nil +} diff --git a/src/framework/utils/ssh/ssh.go b/src/framework/utils/ssh/ssh.go index eaf5c3d9..29a5fb70 100644 --- a/src/framework/utils/ssh/ssh.go +++ b/src/framework/utils/ssh/ssh.go @@ -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 diff --git a/src/framework/utils/ssh/ssh_session.go b/src/framework/utils/ssh/ssh_session.go new file mode 100644 index 00000000..1244ac2b --- /dev/null +++ b/src/framework/utils/ssh/ssh_session.go @@ -0,0 +1,73 @@ +package ssh + +import ( + "bytes" + "fmt" + "io" + "sync" + + gossh "golang.org/x/crypto/ssh" +) + +// 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() +}