codename-sendaround/cmd/sendaround/main.go

package main

import (
	"io"
	"log"
	"os"
	"path/filepath"

	"github.com/cheggaaa/pb/v3"
	"github.com/icedream/sendaround"

	kingpin "gopkg.in/alecthomas/kingpin.v2"
)

var (
	cli           = kingpin.New("sendaround", "CLI to send and receive file transfers via the Sendaround service.")
	flagServerURL = cli.Flag("url", "The Sendaround server URL.").
			URL()
	cmdRecv = cli.Command("receive", "Receive a file transfer using a token.").
		Alias("recv").
		Alias("r")
	cmdRecvArgToken = cmdRecv.Arg("token", "The token of the session to connect to. You receive this token directly from the sender or as part of a URL.").
			Required().
			String()
	cmdRecvFlagOutputDirectory = cmdRecv.Flag("output-directory", "The directory to store the downloaded files to. Defaults to current directory.").
					Short('o').
					Default(".").
					ExistingDir()
	cmdSend = cli.Command("send", "Send a file.").
		Alias("s")
	cmdSendArgFiles = cmdSend.Arg("input-file", "The file to be transmitted.").
			Required().
			ExistingFiles()
)

func main() {
	subcmd := kingpin.MustParse(cli.Parse(os.Args[1:]))

	cfg := new(sendaround.SendaroundClientConfiguration)
	if flagServerURL != nil {
		cfg.ServerURL = *flagServerURL
	}
	c := sendaround.NewSendaroundClient(cfg)

	switch subcmd {
	case cmdRecv.FullCommand():
		conn, err := c.Receive(*cmdRecvArgToken)
		if err != nil {
			log.Fatal(err)
			return
		}
		defer conn.Close()

	connectLoop:
		for state := range conn.StateC() {
			switch state.Type {
			case sendaround.Failed:
				log.Fatal("Failed:", state.Error)
				return
			case sendaround.Disconnected:
				log.Fatal("Early disconnect, aborting.")
				return
			case sendaround.Connected:
				break connectLoop
			}
		}

		pbars := map[string]*pb.ProgressBar{}
		files := conn.Files()
		for _, f := range files {
			bar := pb.New64(int64(f.Length()))
			bar.Set(pb.Bytes, true)
			bar.Set("prefix", f.FileName()+": ")
			pbars[f.FileName()] = bar
		}

		go func() {
			for state := range conn.StateC() {
				switch state.Type {
				case sendaround.TransmittingFile:
					if !pbars[state.CurrentFile.FileName()].IsStarted() {
						pbars[state.CurrentFile.FileName()].Start()
					}
					pbars[state.CurrentFile.FileName()].SetCurrent(int64(state.TransmittedLength))
				case sendaround.Connected:
					if state.CurrentFile != nil { // Was a file transferred successfully just now?
						if state.TransmittedLength == state.CurrentFile.Length() {
							pbars[state.CurrentFile.FileName()].Finish()
						}
					}
				case sendaround.Failed:
					log.Fatal(state.Error)
					return
				}
			}
		}()

		for filePath, f := range files {
			r, err := conn.RetrieveFile(filePath)
			if err != nil {
				log.Fatal(err)
				return
			}
			ofPath := filepath.Join(*cmdRecvFlagOutputDirectory, filepath.FromSlash(f.FileName()))
			of, err := os.Create(ofPath)
			if err != nil {
				log.Fatal(err)
				return
			}
			defer of.Close()
			_, err = io.Copy(of, r)
			if err != nil {
				log.Fatal(err)
				return
			}
		}

	case cmdSend.FullCommand():
		log.Println("Preparing offer...")
		conn, err := c.Offer()
		if err != nil {
			log.Fatal(err)
			return
		}
		defer conn.Close()

		log.Println("")
		log.Println("=============================")
		log.Printf("Token: %s", conn.Token())
		log.Printf("URL: %s", conn.URL())
		log.Println("=============================")
		log.Println("")

		pbars := map[string]*pb.ProgressBar{}

		for _, inputFilePath := range *cmdSendArgFiles {
			f, err := os.Open(inputFilePath)
			if err != nil {
				log.Fatal(err)
				return
			}
			rf, err := sendaround.FileFromFilesystem(f, "")
			if err != nil {
				log.Fatal(err)
				return
			}
			err = conn.AddFile(rf)
			if err != nil {
				log.Fatal(err)
				return
			}
			log.Printf("%s will be transmitted as %s", f.Name(), rf.FileName())
			bar := pb.New64(int64(rf.Length()))
			bar.Set(pb.Bytes, true)
			bar.Set("prefix", rf.FileName()+": ")
			pbars[rf.FileName()] = bar
		}

		log.Println("Now waiting for client...")
		for state := range conn.StateC() {
			switch state.Type {
			case sendaround.TransmittingFile:
				if !pbars[state.CurrentFile.FileName()].IsStarted() {
					pbars[state.CurrentFile.FileName()].Start()
				}
				pbars[state.CurrentFile.FileName()].SetCurrent(int64(state.TransmittedLength))
			case sendaround.Connected:
				if state.CurrentFile != nil { // Was a file transferred successfully just now?
					if state.TransmittedLength == state.CurrentFile.Length() {
						pbars[state.CurrentFile.FileName()].Finish()
					}
				}
			case sendaround.Failed:
				log.Fatal("Failed:", state.Error)
				return
			}
		}
	}
}