Go : ssh操作(三)上传

简介

上传本地文件到远程主机

代码

uploader.go

package client
 
import (
	"github.com/pkg/sftp"
	"io"
	"io/ioutil"
	"os"
	"path"
)
 
var oneTimeMaxSizeToWrite = 8192 // 单次最大写文件大小
 
func IsDir(path string) bool {
	info, err := os.Stat(path)
	if err != nil {
		return false
	}
	return info.IsDir()
}
 
func IsFile(path string) bool {
	info, err := os.Stat(path)
	if err != nil {
		return false
	}
	return !info.IsDir()
}
 
type Uploader interface {
	Upload(localPath, remotePath string) error
	GetUProcess() (num, downloaded uint)
 
	UploadWithCallback(localPath, remotePath string,
		process func(num, downloaded uint),
		done func(err error))
}
 
type uploader struct {
	sftpClient *sftp.Client
 
	uploadSize uint
	upNumber   uint
	uploaded   uint
}
 
func (u *uploader) uploadFolderCount(localPath string) (needUpload, size uint, err error) {
	infos, err := ioutil.ReadDir(localPath)
 
	for _, info := range infos {
		if info.IsDir() {
			c, s, err := u.uploadFolderCount(path.Join(localPath, info.Name()))
			if nil != err { return  }
			needUpload += c
			size += s
			continue
		}
		needUpload += 1
		size += uint(info.Size())
	}
	err = nil
	return
}
 
func (u *uploader) uploadFileCount(localpath string) (uint, uint, error) {
	var (
		isExist bool
		isDir   bool
	)
	info, err := os.Stat(localpath)
	if err != nil {
		isExist = !os.IsNotExist(err)
		isDir = false
	} else {
		isExist = true
		isDir = info.IsDir()
	}
 
	if !isExist { return 0, 0, nil }
	if !isDir { return 1, uint(info.Size()), nil }
 
	return u.uploadFolderCount(localpath)
}
 
func (u *uploader) GetUProcess() (num, uploaded uint) {
	return u.upNumber, u.uploaded
}
 
func (u *uploader) Upload(localPath, remotePath string) (err error) {
	return u.upload(localPath, remotePath, nil)
}
 
func (u *uploader) UploadWithCallback(localPath, remotePath string,
	process func(num, downloaded uint),
	done func(err error)) {
	var err = u.upload(localPath, remotePath, process)
	go done(err)
}
 
func (u *uploader) upload(localPath, remotePath string, process func(num, downloaded uint)) (err error) {
	u.upNumber, u.uploadSize, err = u.uploadFileCount(localPath)
	if nil != err { return err }
 
	var isDir = IsDir(localPath)
	if isDir {
		return u.uploadFolder(localPath, remotePath, process)
	}
 
	return u.uploadFile(localPath, remotePath, process)
}
 
func (u *uploader) writeFile(reader io.Reader, writer io.Writer) (err error) {
	var buffer = make([]byte, oneTimeMaxSizeToWrite)
	var n int
	for ; ; {
		n, err = reader.Read(buffer)
		if n < oneTimeMaxSizeToWrite {
			if io.EOF == err {
				err = nil
				if n > 0 {
					_, err = writer.Write(buffer[0:n])
					if err != nil { return err }
				}
				break
			}
		}
		_, err = writer.Write(buffer)
		if err != nil { return err }
	}
	return nil
}
 
func (u *uploader) uploadFile(localPath, remotePath string, process func(num, downloaded uint)) error {
	var (
		srcFile *os.File
		dstFile *sftp.File
		remoteFileName string
		err error
	)
	srcFile, err = os.Open(localPath)
	if err != nil {
		return err
	}
	defer srcFile.Close()
 
	remoteFileName = path.Base(localPath)
	dstFile, err = u.sftpClient.Create(path.Join(remotePath, remoteFileName))
	if err != nil {
		return err
	}
	defer dstFile.Close()
 
	err = u.writeFile(srcFile, dstFile)
	if nil != err { return err }
 
	u.uploaded += 1
	if nil != process {
		go process(u.upNumber, u.uploaded)
	}
	return nil
}
 
func (u *uploader) uploadFolder(localPath, remotePath string, process func(num, downloaded uint)) (err error) {
	err = u.sftpClient.MkdirAll(remotePath)
	if nil != err { return err }
 
	localFileInfos, err := ioutil.ReadDir(localPath)
	if err != nil {
		return err
	}
 
	for _, fileInfo := range localFileInfos {
		localFilePath := path.Join(localPath, fileInfo.Name())
 
		if fileInfo.IsDir() {
			remoteFilePath := path.Join(remotePath, fileInfo.Name())
 
			err = u.uploadFolder(localFilePath, remoteFilePath, process)
			if nil != err { return err }
		} else {
			err = u.uploadFile(localFilePath, remotePath, process)
			if nil != err { return err }
		}
	}
 
	return nil
}

测试用例

config := ssh.ClientConfig{
		User: "root", // 用户名
		Auth: []ssh.AuthMethod{ssh.Password("xxxx")}, // 密码
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			return nil
		},
		Timeout: 10 * time.Second,
	}
	sshClient, err := ssh.Dial("tcp",
		"192.168.31.75:22", &config) //IP + 端口
	if  err != nil {
		fmt.Print(err)
		return
	}
	defer sshClient.Close()
 
	var sftpClient *sftp.Client
	if sftpClient, err = sftp.NewClient(sshClient); err != nil {
		return
	}
	defer sftpClient.Close()
 
	var up = uploader{
		sftpClient: sftpClient,
	}
 
	up.Upload("C:/test", "/root/")
 
	up.UploadWithCallback("C:/test", "/root/", func(num, down uint) {
		fmt.Println("进度变化: ", num, down)
	}, func (err error) {
		fmt.Println("结束: ", err)
	})
	time.Sleep(3000)