gitee.com/h79/goutils@v1.22.10/common/ssh/scp.go

package ssh

import (
	"gitee.com/h79/goutils/common/file"
	"gitee.com/h79/goutils/common/system"
	"go.uber.org/zap"
	"golang.org/x/crypto/ssh"
	"io"
	"strings"
	"sync"
	"time"
)

const KRemoteCmd = "scp"

type Scp struct {
	session *Session
	cmd     string
}

type Option struct {
	Cmd              string
	Local            Path
	Remote           Path
	UpdatePermission bool //-p 保留原文件的修改时间，访问时间和访问权限。
	Recursive        bool //-r
}

func NewScp(session *Session) *Scp {
	return &Scp{session: session}
}

func (scp *Scp) WithCmd(cmd string) *Scp {
	scp.cmd = cmd
	return scp
}

func (scp *Scp) connect() (*ssh.Session, error) {
	scp.close()
	if err := scp.session.Connect(); err != nil {
		return nil, err
	}
	return scp.session.Session, nil
}

func (scp *Scp) Close() {
	scp.close()
}

func (scp *Scp) close() {
	scp.session.Close()
}

// SendTo 上传
func (scp *Scp) SendTo(id int, opt *Option) *Result {
	dir := file.IsDir(opt.Local.Name)
	if dir == 0 {
		opt.Local.IsDir = false
		//文件名
		return scp.sendFile(id, opt)
	}
	if dir == 1 {
		opt.Local.IsDir = true
		if opt.Remote.IsDir {
			//Remote is dir
			return scp.sendDir(id, opt)
		} else {
			// Remote is file
			localFile(opt)
			return scp.sendFile(id, opt)
		}
	}
	return nil
}

// 发送单文件
func (scp *Scp) sendFile(id int, opt *Option) *Result {
	result := &Result{
		Id:         id,
		Host:       scp.session.Host,
		LocalPath:  opt.Local.Name,
		RemotePath: opt.Remote.Name,
	}
	defer func() {
		result.EndTime = time.Now()
	}()
	start := time.Now()
	result.StartTime = start

	src, size, err := file.Open(opt.Local.Name)
	if err != nil {
		result.Error = err
		return result
	}
	defer src.Close()

	result.Error = scp.SendToByReader(src, size, opt, remoteFile(opt))

	return result
}

// 发送路径文件
func (scp *Scp) sendDir(id int, opt *Option) *Result {
	return nil
}

func (scp *Scp) SendToByReader(src io.Reader, size int64, opt *Option, remoteFilename string) error {
	scp.send(opt)
	return scp.handler(func(session *ssh.Session, out io.Reader, w io.WriteCloser) error {
		var (
			err error
		)
		wg := sync.WaitGroup{}
		wg.Add(2)
		go func() {
			defer system.Recover()
			defer wg.Done()
			defer w.Close()

			sender := &Sender{w: w, info: FileInfo{
				Name: remoteFilename,
				Size: size,
				Mode: opt.Remote.Mode,
			}}
			err = sender.Do(out, src)
		}()

		go func() {
			defer system.Recover()
			defer wg.Done()
			err = session.Run(scp.cmd)
			if err != nil {
				return
			}
		}()
		wg.Wait()
		return err
	})
}

func (scp *Scp) ReceiveFrom(opt *Option) error {

	scp.recv(opt)
	return scp.handler(func(session *ssh.Session, out io.Reader, in io.WriteCloser) error {
		var (
			err error
			wg  = sync.WaitGroup{}
		)
		wg.Add(1)
		go func() {
			defer system.Recover()
			defer wg.Done()
			defer in.Close()

			if err = session.Start(scp.cmd); err != nil {
				return
			}
			receiver := &Receiver{
				opt: opt,
			}
			if err = receiver.Do(out, in); err != nil {
				return
			}
			err = session.Wait()
		}()
		wg.Wait()
		return err
	})
}

func (scp *Scp) handler(start func(session *ssh.Session, out io.Reader, in io.WriteCloser) error) error {
	session, err := scp.connect()
	if err != nil {
		return err
	}
	out, err := session.StdoutPipe()
	if err != nil {
		return err
	}
	in, err := session.StdinPipe()
	if err != nil {
		return err
	}

	return start(session, out, in)
}

func (scp *Scp) recv(opt *Option) {

	if opt.Cmd == "" {
		opt.Cmd = KRemoteCmd
	}

	p := []byte("-f")
	if opt.UpdatePermission {
		p = append(p, 'p')
	}
	if opt.Recursive {
		p = append(p, 'r')
	}
	if opt.Remote.IsDir {
		p = append(p, 'd')
	}
	cmd := opt.Cmd + " " + string(p) + " " + escapeArg(opt.Remote.Name)

	scp.cmd = cmd

	zap.L().Debug("Scp:recv", zap.String("Cmd", cmd))
}

func (scp *Scp) send(opt *Option) {

	if opt.Cmd == "" {
		opt.Cmd = KRemoteCmd
	}

	p := []byte("-qt")
	if opt.UpdatePermission {
		p = append(p, 'p')
	}
	if opt.Recursive {
		p = append(p, 'r')
	}
	if opt.Local.IsDir {
		p = append(p, 'd')
	}

	cmd := opt.Cmd + " " + string(p) + " " + escapeArg(opt.Remote.Name)

	scp.cmd = cmd

	zap.L().Debug("Scp: send", zap.String("Cmd", cmd))
}

func escapeArg(arg string) string {
	return "'" + strings.Replace(arg, "'", `'\''`, -1) + "'"
}