scp: add progress bar

cmd/scpw/main.go

@@ -94,6 +94,8 @@ func Run(ctx *cli.Context) error {
 }
 
 func initScpCli(ctx *cli.Context, node *scpw.Node) error {
+	p := scpw.NewProgress()
+	keepTime := ctx.Bool("keep-time")
 	wg := sync.WaitGroup{}
 	todo := make(chan scpw.LRMap, 5)
 	for i := 0; i < threads; i++ {
@@ -103,15 +105,16 @@ func initScpCli(ctx *cli.Context, node *scpw.Node) error {
 			if err != nil {
 				panic(err)
 			}
-			keepTime := ctx.Bool("keep-time")
 			scpwCli := scpw.NewSCP(ssh, keepTime)
 			defer func() {
 				ssh.Close()
 				wg.Done()
 			}()
 			for lr := range todo {
 				local, remote := lr.Local, lr.Remote
-				err = scpwCli.SwitchScpwFunc(ctx.Context, local, remote, node.Typ)
+				scpwCtx := scpw.Context{Ctx: ctx.Context, Bar: p.NewInfiniteByesBar(local)}
+				err = scpwCli.SwitchScpwFunc(scpwCtx, local, remote, node.Typ)
+				scpwCtx.Bar.SetTotal(-1, true)
 				if err != nil {
 					panic(err)
 				}
@@ -123,5 +126,6 @@ func initScpCli(ctx *cli.Context, node *scpw.Node) error {
 	}
 	close(todo)
 	wg.Wait()
+	p.Wait()
 	return nil
 }

go.mod

@@ -3,26 +3,29 @@ module github.com/T-TRz879/scpw
 go 1.19
 
 require (
-	github.com/go-cmd/cmd v1.4.1
+	github.com/google/gops v0.3.27
 	github.com/google/uuid v1.3.0
 	github.com/manifoldco/promptui v0.9.0
-	github.com/mattn/go-isatty v0.0.14
 	github.com/sirupsen/logrus v1.9.0
 	github.com/stretchr/testify v1.8.1
 	github.com/urfave/cli/v2 v2.4.0
+	github.com/vbauerster/mpb/v8 v8.7.2
 	golang.org/x/crypto v0.4.0
 	gopkg.in/yaml.v2 v2.4.0
-	github.com/google/gops v0.3.27
 )
 
 require (
+	github.com/VividCortex/ewma v1.2.0 // indirect
+	github.com/acarl005/stripansi v0.0.0-20180116102854-5a71ef0e047d // indirect
 	github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e // indirect
 	github.com/cpuguy83/go-md2man/v2 v2.0.2 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/kr/pretty v0.3.0 // indirect
+	github.com/mattn/go-runewidth v0.0.15 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/rivo/uniseg v0.4.4 // indirect
 	github.com/russross/blackfriday/v2 v2.1.0 // indirect
-	golang.org/x/sys v0.5.0 // indirect
+	golang.org/x/sys v0.16.0 // indirect
 	gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -1,4 +1,8 @@
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
+github.com/VividCortex/ewma v1.2.0 h1:f58SaIzcDXrSy3kWaHNvuJgJ3Nmz59Zji6XoJR/q1ow=
+github.com/VividCortex/ewma v1.2.0/go.mod h1:nz4BbCtbLyFDeC9SUHbtcT5644juEuWfUAUnGx7j5l4=
+github.com/acarl005/stripansi v0.0.0-20180116102854-5a71ef0e047d h1:licZJFw2RwpHMqeKTCYkitsPqHNxTmd4SNR5r94FGM8=
+github.com/acarl005/stripansi v0.0.0-20180116102854-5a71ef0e047d/go.mod h1:asat636LX7Bqt5lYEZ27JNDcqxfjdBQuJ/MM4CN/Lzo=
 github.com/chzyer/logex v1.1.10 h1:Swpa1K6QvQznwJRcfTfQJmTE72DqScAa40E+fbHEXEE=
 github.com/chzyer/logex v1.1.10/go.mod h1:+Ywpsq7O8HXn0nuIou7OrIPyXbp3wmkHB+jjWRnGsAI=
 github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e h1:fY5BOSpyZCqRo5OhCuC+XN+r/bBCmeuuJtjz+bCNIf8=
@@ -12,9 +16,6 @@ github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ3
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/go-cmd/cmd v1.4.1 h1:JUcEIE84v8DSy02XTZpUDeGKExk2oW3DA10hTjbQwmc=
-github.com/go-cmd/cmd v1.4.1/go.mod h1:tbBenttXtZU4c5djS1o7PWL5pd2xAr5sIqH1kGdNiRc=
-github.com/go-test/deep v1.0.7 h1:/VSMRlnY/JSyqxQUzQLKVMAskpY/NZKFA5j2P+0pP2M=
 github.com/google/gops v0.3.27 h1:BDdWfedShsBbeatZ820oA4DbVOC8yJ4NI8xAlDFWfgI=
 github.com/google/gops v0.3.27/go.mod h1:lYqabmfnq4Q6UumWNx96Hjup5BDAVc8zmfIy0SkNCSk=
 github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
@@ -28,10 +29,13 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
 github.com/manifoldco/promptui v0.9.0 h1:3V4HzJk1TtXW1MTZMP7mdlwbBpIinw3HztaIlYthEiA=
 github.com/manifoldco/promptui v0.9.0/go.mod h1:ka04sppxSGFAtxX0qhlYQjISsg9mR4GWtQEhdbn6Pgg=
-github.com/mattn/go-isatty v0.0.14 h1:yVuAays6BHfxijgZPzw+3Zlu5yQgKGP2/hcQbHb7S9Y=
-github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94=
+github.com/mattn/go-runewidth v0.0.15 h1:UNAjwbU9l54TA3KzvqLGxwWjHmMgBUVhBiTjelZgg3U=
+github.com/mattn/go-runewidth v0.0.15/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
+github.com/rivo/uniseg v0.4.4 h1:8TfxU8dW6PdqD27gjM8MVNuicgxIjxpm4K7x4jp8sis=
+github.com/rivo/uniseg v0.4.4/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
 github.com/rogpeppe/go-internal v1.6.1 h1:/FiVV8dS/e+YqF2JvO3yXRFbBLTIuSDkuC7aBOAvL+k=
 github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTEfhy4qGm1nDQc=
 github.com/russross/blackfriday/v2 v2.1.0 h1:JIOH55/0cWyOuilr9/qlrm0BSXldqnqwMsf35Ld67mk=
@@ -48,13 +52,14 @@ github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKs
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
 github.com/urfave/cli/v2 v2.4.0 h1:m2pxjjDFgDxSPtO8WSdbndj17Wu2y8vOT86wE/tjr+I=
 github.com/urfave/cli/v2 v2.4.0/go.mod h1:NX9W0zmTvedE5oDoOMs2RTC8RvdK98NTYZE5LbaEYPg=
+github.com/vbauerster/mpb/v8 v8.7.2 h1:SMJtxhNho1MV3OuFgS1DAzhANN1Ejc5Ct+0iSaIkB14=
+github.com/vbauerster/mpb/v8 v8.7.2/go.mod h1:ZFnrjzspgDHoxYLGvxIruiNk73GNTPG4YHgVNpR10VY=
 golang.org/x/crypto v0.4.0 h1:UVQgzMY87xqpKNgb+kDsll2Igd33HszWHFLmpaRMq/8=
 golang.org/x/crypto v0.4.0/go.mod h1:3quD/ATkf6oY+rnes5c3ExXTbLc8mueNue5/DoinL80=
 golang.org/x/sys v0.0.0-20181122145206-62eef0e2fa9b/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
-golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.5.0 h1:MUK/U/4lj1t1oPg0HfuXDN/Z1wv31ZJ/YcPiGccS4DU=
-golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.16.0 h1:xWw16ngr6ZMtmxDyKyIgsE93KNKz5HKmMa3b8ALHidU=
+golang.org/x/sys v0.16.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/term v0.3.0 h1:qoo4akIqOcDME5bhc/NgxUdovd6BSS2uMsVjB56q1xI=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=

progress_utils.go

@@ -0,0 +1,29 @@
+package scpw
+
+import (
+	"fmt"
+	"github.com/vbauerster/mpb/v8"
+	"github.com/vbauerster/mpb/v8/decor"
+)
+
+type Progress struct {
+	*mpb.Progress
+	bars []*mpb.Bar
+}
+
+func NewProgress() *Progress {
+	return &Progress{
+		mpb.New(mpb.WithWidth(64)),
+		[]*mpb.Bar{},
+	}
+}
+
+func (p *Progress) NewInfiniteByesBar(name string) *mpb.Bar {
+	// new bar with 'trigger complete event' disabled, because total is zero
+	bar := p.AddBar(0,
+		mpb.PrependDecorators(decor.Counters(decor.SizeB1024(0), fmt.Sprintf("%-35s", name)+" | % .1f / % .1f")),
+		mpb.AppendDecorators(decor.Percentage()),
+	)
+	p.bars = append(p.bars, bar)
+	return bar
+}

progress_utils_test.go

@@ -0,0 +1,12 @@
+package scpw
+
+import (
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestNewProgress(t *testing.T) {
+	progress := NewProgress()
+	bar := progress.NewInfiniteByesBar("")
+	assert.True(t, bar.IsRunning())
+}

scp.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"fmt"
 	"github.com/google/uuid"
+	"github.com/vbauerster/mpb/v8"
 	"golang.org/x/crypto/ssh"
 	"io"
 	"net"
@@ -20,6 +21,8 @@ var log = GetLogger("scpw")
 
 type CommandType = string
 
+var ONCE_LEN = 128 * 1024
+
 const (
 	C    CommandType = "C"
 	D    CommandType = "D"
@@ -28,6 +31,11 @@ const (
 	NULL CommandType = "NULL"
 )
 
+type Context struct {
+	Ctx context.Context
+	Bar *mpb.Bar
+}
+
 type File struct {
 	Name       string
 	LocalPath  string
@@ -147,7 +155,7 @@ func NewSCP(cli *ssh.Client, keepTime bool) *SCP {
 	}
 }
 
-func (scp *SCP) SwitchScpwFunc(ctx context.Context, localPath, remotePath string, typ SCPWType) error {
+func (scp *SCP) SwitchScpwFunc(ctx Context, localPath, remotePath string, typ SCPWType) error {
 	excludeRootDir := false
 	if typ == PUT {
 		if localPath[len(localPath)-1] == '*' {
@@ -215,7 +223,7 @@ func (scp *SCP) replaceDir(tmp, local, remote string) error {
 	return os.RemoveAll(tmp)
 }
 
-func (scp *SCP) PutAllExcludeRoot(ctx context.Context, srcPath, dstPath string) error {
+func (scp *SCP) PutAllExcludeRoot(ctx Context, srcPath, dstPath string) error {
 	var err error
 	child, err := StatDirChild(srcPath)
 	if err != nil {
@@ -235,7 +243,7 @@ func (scp *SCP) PutAllExcludeRoot(ctx context.Context, srcPath, dstPath string)
 	return nil
 }
 
-func (scp *SCP) PutAll(ctx context.Context, srcPath, dstPath string) error {
+func (scp *SCP) PutAll(ctx Context, srcPath, dstPath string) error {
 	wg := sync.WaitGroup{}
 	wg.Add(2)
 	session, err := scp.NewSession()
@@ -305,7 +313,7 @@ func (scp *SCP) PutAll(ctx context.Context, srcPath, dstPath string) error {
 						return
 					}
 					open, e := os.Open(file.LocalPath)
-					e = parseContent(stdin, open, sizeNum)
+					e = parseContent(ctx.Bar, stdin, open, sizeNum)
 					open.Close()
 					if e != nil {
 						errChan <- e
@@ -322,7 +330,7 @@ func (scp *SCP) PutAll(ctx context.Context, srcPath, dstPath string) error {
 						return
 					}
 				}
-				fmt.Printf("    file:[%40s] size:[%15s]\n", file.Name, file.Size)
+				//fmt.Printf("    file:[%40s] size:[%15s]\n", file.Name, file.Size)
 			case <-scpCh.exitChan:
 				_, err = fmt.Fprintln(stdin, E)
 				if err != nil {
@@ -364,10 +372,10 @@ func (scp *SCP) PutAll(ctx context.Context, srcPath, dstPath string) error {
 	return nil
 }
 
-func WalkTree(ctx context.Context, scpChan *scpChan, rootParent, root, dstPath string) error {
+func WalkTree(ctx Context, scpChan *scpChan, rootParent, root, dstPath string) error {
 	select {
-	case <-ctx.Done():
-		return ctx.Err()
+	case <-ctx.Ctx.Done():
+		return ctx.Ctx.Err()
 	default:
 		child, name, mode, atime, mtime, _, err := StatDir(root)
 		if err != nil {
@@ -404,7 +412,7 @@ func WalkTree(ctx context.Context, scpChan *scpChan, rootParent, root, dstPath s
 	}
 }
 
-func (scp *SCP) Put(ctx context.Context, srcPath, dstPath string) error {
+func (scp *SCP) Put(ctx Context, srcPath, dstPath string) error {
 	stat, err := os.Stat(srcPath)
 	if err != nil {
 		return err
@@ -427,7 +435,7 @@ func (scp *SCP) Put(ctx context.Context, srcPath, dstPath string) error {
 	return scp.put(ctx, dstPath, open, mode, stat.Size(), atime, mtime)
 }
 
-func (scp *SCP) put(ctx context.Context, dstPath string, in io.Reader, mode string, size int64, atime, mtime string) error {
+func (scp *SCP) put(ctx Context, dstPath string, in io.Reader, mode string, size int64, atime, mtime string) error {
 	wg := sync.WaitGroup{}
 	session, err := scp.NewSession()
 	if err != nil {
@@ -496,7 +504,7 @@ func (scp *SCP) put(ctx context.Context, dstPath string, in io.Reader, mode stri
 			errChan <- err
 			return
 		}
-		fmt.Printf("    file:[%40s] size:[%15d]\n", fileName, size)
+		//fmt.Printf("    file:[%40s] size:[%15d]\n", fileName, size)
 	}()
 
 	var err2 error
@@ -524,7 +532,7 @@ func (scp *SCP) put(ctx context.Context, dstPath string, in io.Reader, mode stri
 	return nil
 }
 
-func (scp *SCP) Get(ctx context.Context, srcPath, dstPath string) error {
+func (scp *SCP) Get(ctx Context, srcPath, dstPath string) error {
 	session, err := scp.NewSession()
 	defer session.Close()
 	if err != nil {
@@ -612,7 +620,7 @@ func (scp *SCP) Get(ctx context.Context, srcPath, dstPath string) error {
 			}
 		}
 
-		if err = parseContent(in, stdout, attr.Size); err != nil {
+		if err = parseContent(ctx.Bar, in, stdout, attr.Size); err != nil {
 			os.Remove(srcPath)
 			errChan <- err
 			return
@@ -636,7 +644,7 @@ func (scp *SCP) Get(ctx context.Context, srcPath, dstPath string) error {
 			os.Remove(srcPath)
 			return
 		}
-		fmt.Printf("    file:[%40s] size:[%15d]\n", filepath.Base(srcPath), attr.Size)
+		//fmt.Printf("    file:[%40s] size:[%15d]\n", filepath.Base(srcPath), attr.Size)
 	}()
 	wg.Wait()
 	close(errChan)
@@ -648,7 +656,7 @@ func (scp *SCP) Get(ctx context.Context, srcPath, dstPath string) error {
 	return nil
 }
 
-func (scp *SCP) GetAll(ctx context.Context, localPath, remotePath string) error {
+func (scp *SCP) GetAll(ctx Context, localPath, remotePath string) error {
 	session, err := scp.NewSession()
 	if err != nil {
 		return err
@@ -729,15 +737,15 @@ func (scp *SCP) GetAll(ctx context.Context, localPath, remotePath string) error
 					os.Remove(curLocal)
 					errChan <- err
 				}
-				fmt.Printf("    file:[%40s] size:[%15d]\n", attr.Name, attr.Size)
+				//fmt.Printf("    file:[%40s] size:[%15d]\n", attr.Name, attr.Size)
 
 			} else if attr.Typ == D {
 				// mkdir dir
 				if e = os.Mkdir(curLocal, attr.Mode); e != nil {
 					errChan <- e
 					return
 				}
-				fmt.Printf("    file:[%40s] size:[%15d]\n", attr.Name, attr.Size)
+				//fmt.Printf("    file:[%40s] size:[%15d]\n", attr.Name, attr.Size)
 
 			} else if attr.Typ == E {
 				// cd ../
@@ -757,7 +765,7 @@ func (scp *SCP) GetAll(ctx context.Context, localPath, remotePath string) error
 				}
 			}
 			if attr.Typ == C {
-				if e = parseContent(in, stdout, attr.Size); e != nil {
+				if e = parseContent(ctx.Bar, in, stdout, attr.Size); e != nil {
 					os.Remove(curLocal)
 					errChan <- e
 					return
@@ -870,13 +878,15 @@ func parseCommandType(s string) CommandType {
 	}
 }
 
-func parseContent(in io.Writer, out io.Reader, size int64) error {
+func parseContent(bar *mpb.Bar, in io.Writer, out io.Reader, size int64) error {
 	var read int64
 	for read < size {
 		if readN, err := io.CopyN(in, out, size); err != nil {
 			return err
 		} else {
 			read += readN
+			bar.IncrBy(int(readN))
+			bar.SetTotal(bar.Current()+readN, false)
 		}
 	}
 	return nil