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

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 {
	// 创建本地目录
	err := os.MkdirAll(localDir, 0775)
	if err != nil {
		logger.Errorf("CopyDirRemoteToLocal failed to creating local directory %s: => %s", localDir, err.Error())
		return err
	}

	// 列出远程目录中的文件和子目录
	remoteFiles, err := s.Client.ReadDir(remoteDir)
	if err != nil {
		logger.Errorf("CopyDirRemoteToLocal failed to reading remote directory %s: => %s", remoteDir, err.Error())
		return err
	}

	// 遍历远程文件和子目录并复制到本地
	for _, remoteFile := range remoteFiles {
		remotePath := filepath.ToSlash(filepath.Join(remoteDir, remoteFile.Name()))
		localPath := filepath.ToSlash(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 {
			// 如果是文件，则复制文件内容
			if err := s.CopyFileRemoteToLocal(remotePath, localPath); err != nil {
				logger.Errorf("CopyDirRemoteToLocal failed to opening remote file %s: => %s", remotePath, err.Error())
				continue
			}
		}
	}
	return nil
}

// CopyDirLocalToRemote 复制目录-本地到远程
func (s *SSHClientSFTP) CopyDirLocalToRemote(localDir, remoteDir string) error {
	// 遍历本地目录中的文件和子目录并复制到远程
	err := filepath.Walk(localDir, func(localPath string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// 生成远程路径
		remotePath := filepath.ToSlash(filepath.Join(remoteDir, localPath[len(localDir):]))

		if info.IsDir() {
			// 如果是子目录，则创建远程目录
			if err := s.Client.MkdirAll(remotePath); err != nil {
				logger.Errorf("CopyDirLocalToRemote failed to creating remote directory %s: => %s", remotePath, err.Error())
				return err
			}
		} else {
			// 如果是文件，则复制文件内容
			if err := s.CopyFileLocalToRemote(localPath, remotePath); err != nil {
				logger.Errorf("CopyDirLocalToRemote failed to copying remote file %s: => %s", localPath, err.Error())
				return err
			}
		}

		return nil
	})
	if err != nil {
		logger.Errorf("CopyDirLocalToRemote failed to walking remote directory: => %s", err.Error())
		return err
	}
	return nil
}

// CopyFileRemoteToLocal 复制文件-远程到本地
func (s *SSHClientSFTP) CopyFileRemoteToLocal(remotePath, localPath string) error {
	if err := os.MkdirAll(filepath.Dir(localPath), 0775); err != nil {
		return err
	}

	// 打开远程文件
	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()

	// 创建本地文件
	localFile, err := os.Create(localPath)
	if err != nil {
		logger.Errorf("CopyFileRemoteToLocal failed to creating local file: => %s", err.Error())
		return err
	}
	defer localFile.Close()

	// 复制文件内容
	if _, err = io.Copy(localFile, remoteFile); err != nil {
		logger.Errorf("CopyFileRemoteToLocal failed to copying contents: => %s", err.Error())
		return err
	}
	return nil
}

// CopyFileLocalToRemote 复制文件-本地到远程
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()

	// 创建远程目录
	// if err := s.Client.MkdirAll(filepath.Dir(remotePath)); err != nil {
	// 	logger.Errorf("CopyFileLocalToRemote failed to creating remote directory %s: => %s", remotePath, err.Error())
	// 	return err
	// }

	// 创建远程文件
	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()

	// 复制文件内容
	if _, err = io.Copy(remoteFile, localFile); err != nil {
		logger.Errorf("CopyFileLocalToRemote failed to copying contents: => %s", err.Error())
		return err
	}
	return nil
}