zerone/sync/sync.go

package main

import (
	"fmt"
	"io"
	"os"
	"os/signal"
	"path/filepath"
	"regexp"
	"strings"
	"syscall"
	"time"

	"git.wecise.com/wecise/util/ssh"
	"github.com/spf13/cast"
	ccfg "github.com/wecisecode/util/cfg"
	"github.com/wecisecode/util/filewalker"
	clog "github.com/wecisecode/util/logger"
)

func init() {
	exitChan := make(chan os.Signal, 1)
	signal.Notify(exitChan, os.Interrupt, os.Kill, syscall.SIGHUP, syscall.SIGQUIT, syscall.SIGABRT, syscall.SIGTERM)
	go func() {
		<-exitChan
		os.Exit(1)
	}()
}

var config = ccfg.MConfig()
var logger = clog.New()

// 同步指定文件夹中的内容

func PrintUsage() {
	println(`Usage:
sync <frompath> <topath> [<copytopath> ...]
<frompath>     源文件路径
<topath>       目标文件路径，增量同步
<copytopath>   副本文件路径，<topath>增量同步
`)
}

var option = struct {
	Debug bool
}{
	Debug: false,
}

func main() {
	var frompath string = ""
	var topath string = ""
	var copytopath = []string{}
	for _, arg := range os.Args[1:] {
		if strings.HasPrefix(arg, "-") {
			switch arg {
			case "-debug":
				option.Debug = true
			}
		} else if arg != "" {
			if frompath == "" {
				frompath = arg
			} else if topath == "" {
				topath = arg
			} else {
				copytopath = append(copytopath, arg)
			}
		}
	}
	switch {
	case frompath != "" && topath != "":
		err := sync(frompath, topath, copytopath...)
		if err != nil {
			println(err.Error())
			os.Exit(1)
		}
	default:
		PrintUsage()
		os.Exit(1)
	}
}

var REremotepath = regexp.MustCompile(`^([^:]+):([^@]+)@([^/]+)(/.*)`)
var REabspath = regexp.MustCompile(`^((?:/.*)|(?:[a-zA-Z]:\.*)|(?:[a-zA-Z]:))`)

func sync(frompath, topath string, copytopath ...string) (err error) {
	frompath, e := filepath.Abs(frompath)
	if e != nil {
		err = e
		return
	}
	topath, e = filepath.Abs(topath)
	if e != nil {
		err = e
		return
	}
	fmt.Println("sync from", frompath, "to", topath, "\ncopy to", copytopath)
	for {
		err = sync_once(frompath, topath, copytopath...)
		if err != nil {
			return
		}
		time.Sleep(5 * time.Second)
	}
}

func sync_once(frompath, topath string, copytopath ...string) (err error) {
	fw, e := filewalker.NewFileWalker([]string{frompath}, `^[^\.].*`) // orderby: dirfirst, filefirst, fullpath
	if e != nil {
		return e
	}
	count := 0
	e = fw.List(func(basedir, fpath string) bool {
		if option.Debug {
			fmt.Println("basedir=", basedir, ", fpath=", fpath)
		}
		absbasedir, e := filepath.Abs(basedir)
		if e != nil {
			err = e
			return false
		}
		hiddenpath := fmt.Sprintf(`.*\%c\..*`, os.PathSeparator)
		if regexp.MustCompile(`^[^\.].*`).MatchString(fpath) && !regexp.MustCompile(hiddenpath).MatchString(fpath) {
			fromfile := filepath.Join(absbasedir, fpath)
			tofile := filepath.Join(topath, fpath)
			if option.Debug {
				fmt.Println(time.Now().Format("2006-01-02 15:04:05"), fromfile, " ==> ", tofile)
			}
			fromfi, e := os.Stat(fromfile)
			if e != nil {
				err = fmt.Errorf("from file %v", e)
				return false
			}
			tofi, e := os.Stat(tofile)
			if e != nil && !os.IsNotExist(e) {
				err = fmt.Errorf("to file %v", e)
				return false
			}
			if tofi != nil && fromfi.Size() == tofi.Size() && fromfi.ModTime() == tofi.ModTime() {
				return true
			}
			progressinfo := fmt.Sprint(fromfile, " ==> ", tofile)
			fmt.Print(progressinfo)
			frombs, e := ReadFile(fromfile, progressinfo)
			if e != nil {
				err = fmt.Errorf("from file %v", e)
				return false
			}
			WriteFile(topath, fpath, frombs, fromfi.ModTime())
			for _, ctpath := range copytopath {
				WriteFile(ctpath, fpath, frombs, fromfi.ModTime())
			}
			count++
		}
		return true
	})
	if e != nil {
		return e
	}
	if err != nil {
		return
	}
	if count > 0 || option.Debug {
		fmt.Println(time.Now().Format("2006-01-02 15:04:05"), count, "files copied")
	}
	return
}

func ReadFile(name string, progressinfo string) (data []byte, err error) {
	defer fmt.Println()
	data = make([]byte, 0, 512)
	for {
		data, err = ReadFileContinue(name, progressinfo, data)
		if err == io.EOF {
			err = nil
			return
		}
	}
}

func ReadFileContinue(name string, progressinfo string, data []byte) ([]byte, error) {
	f, err := os.Open(name)
	if err != nil {
		return data, err
	}
	defer f.Close()

	var filesize int64
	var bufsize int
	if info, err := f.Stat(); err == nil {
		filesize = info.Size()
		if int64(int(filesize)) == filesize {
			bufsize = int(filesize)
		}
	}
	bufsize++ // one byte for final read at EOF

	// If a file claims a small size, read at least 512 bytes.
	// In particular, files in Linux's /proc claim size 0 but
	// then do not work right if read in small pieces,
	// so an initial read of 1 byte would not work correctly.
	offset := int64(len(data))
	if bufsize > cap(data) {
		ndata := make([]byte, len(data), bufsize)
		copy(ndata, data)
		data = ndata
	}

	statuschan := make(chan int64, 100)
	done := make(chan struct{}, 1)
	go func() {
		rateinfo := ""
		for range statuschan {
			// bs := strings.Repeat("\b", len(progressinfo+rateinfo))
			rateinfo := fmt.Sprintf("    %.2f%%    ", 100*float64(offset)/float64(filesize))
			// fmt.Printf("\r%s%s", bs, progressinfo+rateinfo)
			fmt.Printf("\r%s", progressinfo+rateinfo)
		}
		if offset >= filesize {
			// bs := strings.Repeat("\b", len(progressinfo+rateinfo))
			rateinfo = fmt.Sprintf("    %.2f%%    ", 100.)
			// fmt.Printf("\r%s%s", bs, progressinfo+rateinfo)
			fmt.Printf("\r%s", progressinfo+rateinfo)
		}
		done <- struct{}{}
	}()
	chunk := int64(1024 * 256)
	t := time.Now()
	for {
		to := offset + chunk
		if to > int64(cap(data)) {
			to = int64(cap(data))
		}
		n, err := f.ReadAt(data[offset:to], int64(offset))
		offset += int64(n)
		statuschan <- offset
		if err != nil {
			close(statuschan)
			<-done
			return data[:offset], err
		}
		if offset >= int64(cap(data)) {
			d := append(data[:cap(data)], 0)
			data = d[:offset]
		}
		ut := time.Since(t)
		if ut < 500*time.Millisecond {
			chunk *= 2
		} else if chunk > 1024 && ut > 5*time.Second {
			chunk /= 2
		}
		t = time.Now()
	}
}

func WriteFile(basedir string, fpath string, bs []byte, mtime time.Time) (err error) {
	ss := REremotepath.FindStringSubmatch(basedir)
	if len(ss) == 5 {
		username, password, hostport, remotepath := ss[1], ss[2], ss[3], regexp.MustCompile("//+").ReplaceAllString(ss[4], "/")
		return WriteRemoteFile(username, password, hostport, remotepath, fpath, bs, mtime)
	}
	tofile := filepath.Join(basedir, fpath)
	todir := filepath.Dir(tofile)
	e := os.MkdirAll(todir, os.ModePerm)
	if e != nil {
		err = fmt.Errorf("to dir %v", e)
		return
	}
	e = os.WriteFile(tofile, bs, os.ModePerm)
	if e != nil {
		err = fmt.Errorf("to file %v", e)
		return
	}
	e = os.Chtimes(tofile, mtime, mtime)
	if e != nil {
		err = fmt.Errorf("to file %v", e)
		return
	}
	return nil
}

func WriteRemoteFile(username, password, hostport, remotepath, fpath string, bs []byte, mtime time.Time) error {
	fpath = strings.ReplaceAll(fpath, "\\", "/")
	fmt.Println("copy to:", fmt.Sprint(username, ":", password, "@", hostport, "/", remotepath, "/", fpath))
	hp := strings.Split(hostport, ":")
	host := hp[0]
	port := 22
	if len(hp) >= 2 && hp[1] != "" {
		port = cast.ToInt(hp[1])
	}
	node := &ssh.Node{
		User:     username,
		Password: password,
		Host:     host,
		Port:     port,
	}
	client, e := node.Connect()
	if e != nil {
		return e
	}
	rfpath := remotepath + "/" + fpath
	diridx := strings.LastIndex(rfpath, "/")
	e = client.MkdirAll(rfpath[:diridx])
	if e != nil {
		return e
	}
	e = client.WriteFile(rfpath, bs, os.ModePerm, mtime)
	if e != nil {
		return e
	}
	return nil
}