gosync/tools/rsync.go

package tools

import (
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
	"velvettear/gosync/log"
	"velvettear/gosync/settings"

	"github.com/fatih/color"
	"github.com/vbauerster/mpb/v8"
	"github.com/vbauerster/mpb/v8/decor"
)

var transferSize float64

// exported function(s)
func Transfer() {
	transferSize = 0
	sourcefiles := getSourceFiles()
	sourcefilesCount := len(sourcefiles)
	if sourcefilesCount <= 0 {
		log.Info("nothing to do - exiting...")
		os.Exit(0)
	}
	var waitgroup sync.WaitGroup
	waitgroup.Add(sourcefilesCount)
	barcontainer := mpb.New(
		mpb.WithWaitGroup(&waitgroup),
	)
	timestamp := time.Now()
	counter := 0
	totalbar := createProgressBar(barcontainer, "Total", int64(0), int64(sourcefilesCount), sourcefilesCount+1, true)
	concurrency := settings.Concurrency
	if concurrency == 0 {
		concurrency = sourcefilesCount
	}
	channel := make(chan struct{}, concurrency)
	for index, file := range sourcefiles {
		channel <- struct{}{}
		if index > 0 {
			time.Sleep(100 * time.Millisecond)
		}
		go func(index int, file string) {
			defer waitgroup.Done()
			stats, error := os.Stat(file)
			if error != nil {
				log.Warning("encountered an error getting file size", error.Error())
				return
			}
			bar := createProgressBar(barcontainer, filepath.Base(file), stats.Size(), int64(100), index, false)
			if transferFile(bar, file) {
				counter++
			}
			totalbar.Increment()
			<-channel
		}(index, file)
	}
	barcontainer.Wait()
	timeDifference := time.Since(timestamp)
	transferSpeedName := "bytes/sec"
	transferSpeed := transferSize / timeDifference.Seconds()
	if transferSpeed > 1048576 {
		transferSpeed = transferSpeed / 1048576
		transferSpeedName = "mb/s"
	}
	transferSizeName := "bytes"
	if transferSize > 1048576 {
		transferSize = transferSize / 1048576
		transferSizeName = "mb"
	}
	log.InfoTimed("transferred "+strconv.Itoa(counter)+" files, "+strconv.Itoa(int(transferSize))+" "+transferSizeName+" ("+strconv.FormatFloat(transferSpeed, 'f', 2, 64)+" "+transferSpeedName+")", timestamp.UnixMilli())
}

// unexported function(s)
func transferFile(bar *mpb.Bar, file string) bool {
	target := getTargetLocation(file)
	var arguments []string
	if len(settings.Password) > 0 {
		arguments = append(arguments, "-p", settings.Password)
	}
	arguments = append(arguments, "rsync", "-avz", "--mkpath", file)
	if len(settings.User) > 0 {
		target = settings.User + "@" + target
	}
	arguments = append(arguments, target, "--progress")
	cmd := exec.Command("sshpass", arguments...)
	stdout, stdoutError := cmd.StdoutPipe()
	stderr, stderrError := cmd.StderrPipe()
	cmd.Start()
	if stdoutError != nil {
		log.Fatal(stdoutError.Error())
	}
	if stderrError != nil {
		log.Fatal(stderrError.Error())
	}
	var resultBytes []byte
	var waitgroup sync.WaitGroup
	waitgroup.Add(1)
	go func() {
		defer waitgroup.Done()
		for {
			tmp := make([]byte, 1024)
			readBytes, error := stdout.Read(tmp)
			if error != nil {
				if error != io.EOF {
					log.Warning("encountered an error reading stdout", error.Error())
				}
				break
			}
			if readBytes == 0 {
				break
			}
			resultBytes = append(resultBytes, tmp...)
			line := string(tmp)
			lowerline := strings.ToLower(line)
			if strings.Contains(lowerline, "sent") && strings.Contains(lowerline, "received") && strings.Contains(lowerline, "bytes") {
				_, tmp, _ := strings.Cut(lowerline, "sent")
				tmp, _, _ = strings.Cut(tmp, "bytes")
				tmp = strings.ReplaceAll(strings.TrimSpace(tmp), ".", "")
				bytes, error := strconv.ParseFloat(tmp, 64)
				if error != nil {
					log.Fatal("encountered an error converting the transferred bytes to int", error.Error())
				}
				transferSize += bytes
			}
			if strings.Contains(lowerline, "total size is") && strings.Contains(lowerline, "speedup is") {
				bar.SetCurrent(100)
				break
			}
			cut := strings.Index(line, "%")
			if cut <= 0 {
				continue
			}
			line = line[:cut]
			var percent string
			for index := len(line) - 1; index > 0; index-- {
				char := string(line[index])
				if char == " " {
					break
				}
				percent = char + percent
			}
			value, error := strconv.Atoi(percent)
			if error != nil {
				continue
			}
			bar.SetCurrent(int64(value))
		}
	}()
	errorBytes, stderrError := io.ReadAll(stderr)
	if stderrError != nil {
		log.Fatal(stderrError.Error())
	}
	cmd.Wait()
	error := strings.Trim(string(errorBytes), "\n")
	if len(error) > 0 {
		log.Fatal(error)
	}
	waitgroup.Wait()
	stdout.Close()
	stderr.Close()
	return true
}

func getTargetLocation(sourceFile string) string {
	if sourceFile == settings.Source {
		return filepath.Join(settings.Target, filepath.Base(sourceFile))
	}
	source, _ := strings.CutSuffix(settings.Source, "/*")
	return filepath.Join(settings.Target, strings.Replace(sourceFile, filepath.Dir(source), "", 1))
}

func createProgressBar(barcontainer *mpb.Progress, name string, size int64, max int64, priority int, total bool) *mpb.Bar {
	red, green, magenta, yellow := color.New(color.FgRed), color.New(color.FgGreen), color.New(color.FgMagenta), color.New(color.FgYellow)
	barstyle := mpb.BarStyle().Lbound("").Filler("󰝤").Tip("").Padding(" ").Rbound("")
	defaultBarPrepend := mpb.PrependDecorators(
		decor.Name("[info] > "),
		decor.OnCompleteMeta(
			decor.OnComplete(
				decor.Meta(decor.Name(name, decor.WCSyncSpaceR), colorize(red)), ""+name,
			),
			colorize(green),
		),
	)
	defaultBarAppend := mpb.AppendDecorators(
		decor.OnCompleteMeta(decor.Elapsed(decor.ET_STYLE_GO, decor.WCSyncSpaceR), colorize(yellow)),
		decor.OnComplete(
			decor.Name("|", decor.WCSyncSpaceR), "",
		),
		decor.OnComplete(
			decor.Percentage(decor.WCSyncSpaceR), "",
		),
	)
	totalBarPrepend := mpb.PrependDecorators(
		decor.Name("[info] > "),
		decor.OnCompleteMeta(
			decor.OnComplete(
				decor.Meta(decor.Name(name, decor.WC{W: len(name) + 1, C: decor.DidentRight}), colorize(yellow)), ""+name,
			),
			colorize(magenta),
		),
		decor.CountersNoUnit("%d / %d"),
	)
	if total {
		return barcontainer.New(
			max,
			barstyle,
			mpb.BarPriority(priority),
			mpb.BarFillerClearOnComplete(),
			totalBarPrepend,
			defaultBarAppend,
		)
	} else {
		return barcontainer.New(
			max,
			barstyle,
			mpb.BarPriority(priority),
			mpb.BarFillerClearOnComplete(),
			defaultBarPrepend,
			defaultBarAppend,
		)
	}
}

func colorize(c *color.Color) func(string) string {
	return func(s string) string {
		return c.Sprint(s)
	}
}