From 47ab799c6fc4c8d695b32a294cdffd5c0d903249 Mon Sep 17 00:00:00 2001 From: Zack Scholl Date: Wed, 17 Oct 2018 06:19:01 -0700 Subject: [PATCH] add file stats to state --- src/croc/croc.go | 7 +++- src/croc/recipient.go | 75 +++++++++++++++++++++---------------------- src/croc/sender.go | 43 ++++++++++++------------- 3 files changed, 64 insertions(+), 61 deletions(-) diff --git a/src/croc/croc.go b/src/croc/croc.go index 144115f..db217b7 100644 --- a/src/croc/croc.go +++ b/src/croc/croc.go @@ -5,8 +5,10 @@ import ( "time" "github.com/schollz/croc/src/logger" + "github.com/schollz/croc/src/models" "github.com/schollz/croc/src/relay" "github.com/schollz/croc/src/zipper" + "github.com/schollz/progressbar" ) func init() { @@ -51,7 +53,10 @@ type Croc struct { isLocal bool normalFinish bool - State string + // state variables + StateString string + Bar *progressbar.ProgressBar + FileInfo models.FileStats } // Init will initiate with the default parameters diff --git a/src/croc/recipient.go b/src/croc/recipient.go index 84b2207..66abddf 100644 --- a/src/croc/recipient.go +++ b/src/croc/recipient.go @@ -45,7 +45,6 @@ func (cr *Croc) startRecipient(forceSend int, serverAddress string, tcpPorts []s } func (cr *Croc) receive(forceSend int, serverAddress string, tcpPorts []string, isLocal bool, c *websocket.Conn, codephrase string, noPrompt bool, useStdout bool) (err error) { - var fstats models.FileStats var sessionKey []byte var transferTime time.Duration var hash256 []byte @@ -172,16 +171,16 @@ func (cr *Croc) receive(forceSend int, serverAddress string, tcpPorts []string, if err != nil { return err } - err = json.Unmarshal(decryptedFileData, &fstats) + err = json.Unmarshal(decryptedFileData, &cr.FileInfo) if err != nil { return err } - log.Debugf("got file stats: %+v", fstats) + log.Debugf("got file stats: %+v", cr.FileInfo) // determine if the file is resuming or not - progressFile = fmt.Sprintf("%s.progress", fstats.SentName) + progressFile = fmt.Sprintf("%s.progress", cr.FileInfo.SentName) overwritingOrReceiving := "Receiving" - if utils.Exists(fstats.Name) || utils.Exists(fstats.SentName) { + if utils.Exists(cr.FileInfo.Name) || utils.Exists(cr.FileInfo.SentName) { overwritingOrReceiving = "Overwriting" if utils.Exists(progressFile) { overwritingOrReceiving = "Resume receiving" @@ -205,14 +204,14 @@ func (cr *Croc) receive(forceSend int, serverAddress string, tcpPorts []string, // prompt user about the file fileOrFolder := "file" - if fstats.IsDir { + if cr.FileInfo.IsDir { fileOrFolder = "folder" } fmt.Fprintf(os.Stderr, "\r%s %s (%s) into: %s\n", overwritingOrReceiving, fileOrFolder, - humanize.Bytes(uint64(fstats.Size)), - fstats.Name, + humanize.Bytes(uint64(cr.FileInfo.Size)), + cr.FileInfo.Name, ) if !noPrompt { if "y" != utils.GetInput("ok? (y/N): ") { @@ -225,27 +224,27 @@ func (cr *Croc) receive(forceSend int, serverAddress string, tcpPorts []string, // await file // erase file if overwriting if overwritingOrReceiving == "Overwriting" { - os.Remove(fstats.SentName) + os.Remove(cr.FileInfo.SentName) } var f *os.File - if utils.Exists(fstats.SentName) && resumeFile { + if utils.Exists(cr.FileInfo.SentName) && resumeFile { if !useWebsockets { - f, err = os.OpenFile(fstats.SentName, os.O_WRONLY, 0644) + f, err = os.OpenFile(cr.FileInfo.SentName, os.O_WRONLY, 0644) } else { - f, err = os.OpenFile(fstats.SentName, os.O_APPEND|os.O_WRONLY, 0644) + f, err = os.OpenFile(cr.FileInfo.SentName, os.O_APPEND|os.O_WRONLY, 0644) } if err != nil { log.Error(err) return err } } else { - f, err = os.Create(fstats.SentName) + f, err = os.Create(cr.FileInfo.SentName) if err != nil { log.Error(err) return err } if !useWebsockets { - if err = f.Truncate(fstats.Size); err != nil { + if err = f.Truncate(cr.FileInfo.Size); err != nil { log.Error(err) return err } @@ -262,13 +261,13 @@ func (cr *Croc) receive(forceSend int, serverAddress string, tcpPorts []string, // start the ui for pgoress bytesWritten := 0 fmt.Fprintf(os.Stderr, "\nReceiving (<-%s)...\n", otherIP) - bar := progressbar.NewOptions( - int(fstats.Size), + cr.Bar = progressbar.NewOptions( + int(cr.FileInfo.Size), progressbar.OptionSetRenderBlankState(true), - progressbar.OptionSetBytes(int(fstats.Size)), + progressbar.OptionSetBytes(int(cr.FileInfo.Size)), progressbar.OptionSetWriter(os.Stderr), ) - bar.Add((len(blocks) * blockSize)) + cr.Bar.Add((len(blocks) * blockSize)) finished := make(chan bool) go func(finished chan bool, dataChan chan []byte) (err error) { @@ -288,7 +287,7 @@ func (cr *Croc) receive(forceSend int, serverAddress string, tcpPorts []string, defer fProgress.Close() blocksWritten := 0.0 - blocksToWrite := float64(fstats.Size) + blocksToWrite := float64(cr.FileInfo.Size) if useWebsockets { blocksToWrite = blocksToWrite/float64(models.WEBSOCKET_BUFFER_SIZE/8) - float64(len(blocks)) } else { @@ -304,7 +303,7 @@ func (cr *Croc) receive(forceSend int, serverAddress string, tcpPorts []string, log.Error(err) return err } - decrypted, err := enc.Decrypt(sessionKey, !fstats.IsEncrypted) + decrypted, err := enc.Decrypt(sessionKey, !cr.FileInfo.IsEncrypted) if err != nil { log.Error(err) return err @@ -319,7 +318,7 @@ func (cr *Croc) receive(forceSend int, serverAddress string, tcpPorts []string, } // do decompression - if fstats.IsCompressed && !fstats.IsDir { + if cr.FileInfo.IsCompressed && !cr.FileInfo.IsDir { decrypted = compress.Decompress(decrypted) } @@ -347,9 +346,9 @@ func (cr *Croc) receive(forceSend int, serverAddress string, tcpPorts []string, bytesWritten += n blocksWritten += 1.0 // update the progress bar - bar.Add(n) - if int64(bytesWritten) == fstats.Size || blocksWritten >= blocksToWrite { - log.Debug("finished", int64(bytesWritten), fstats.Size, blocksWritten, blocksToWrite) + cr.Bar.Add(n) + if int64(bytesWritten) == cr.FileInfo.Size || blocksWritten >= blocksToWrite { + log.Debug("finished", int64(bytesWritten), cr.FileInfo.Size, blocksWritten, blocksToWrite) break } } @@ -432,10 +431,10 @@ func (cr *Croc) receive(forceSend int, serverAddress string, tcpPorts []string, } // finish bar - bar.Finish() + cr.Bar.Finish() // check hash - hash256, err = utils.HashFile(fstats.SentName) + hash256, err = utils.HashFile(cr.FileInfo.SentName) if err != nil { log.Error(err) return err @@ -447,45 +446,45 @@ func (cr *Croc) receive(forceSend int, serverAddress string, tcpPorts []string, log.Debugf("got hash: %x", message) if bytes.Equal(hash256, message) { // open directory - if fstats.IsDir { - err = zipper.UnzipFile(fstats.SentName, ".") + if cr.FileInfo.IsDir { + err = zipper.UnzipFile(cr.FileInfo.SentName, ".") if DebugLevel != "debug" { - os.Remove(fstats.SentName) + os.Remove(cr.FileInfo.SentName) } } else { err = nil } if err == nil { - if useStdout && !fstats.IsDir { + if useStdout && !cr.FileInfo.IsDir { var bFile []byte - bFile, err = ioutil.ReadFile(fstats.SentName) + bFile, err = ioutil.ReadFile(cr.FileInfo.SentName) if err != nil { return err } os.Stdout.Write(bFile) - os.Remove(fstats.SentName) + os.Remove(cr.FileInfo.SentName) } - transferRate := float64(fstats.Size) / 1000000.0 / transferTime.Seconds() + transferRate := float64(cr.FileInfo.Size) / 1000000.0 / transferTime.Seconds() transferType := "MB/s" if transferRate < 1 { - transferRate = float64(fstats.Size) / 1000.0 / transferTime.Seconds() + transferRate = float64(cr.FileInfo.Size) / 1000.0 / transferTime.Seconds() transferType = "kB/s" } folderOrFile := "file" - if fstats.IsDir { + if cr.FileInfo.IsDir { folderOrFile = "folder" } if useStdout { - fstats.Name = "stdout" + cr.FileInfo.Name = "stdout" } - fmt.Fprintf(os.Stderr, "\nReceived %s written to %s (%2.1f %s)\n", folderOrFile, fstats.Name, transferRate, transferType) + fmt.Fprintf(os.Stderr, "\nReceived %s written to %s (%2.1f %s)\n", folderOrFile, cr.FileInfo.Name, transferRate, transferType) os.Remove(progressFile) } return err } else { if DebugLevel != "debug" { log.Debug("removing corrupted file") - os.Remove(fstats.SentName) + os.Remove(cr.FileInfo.SentName) } return errors.New("file corrupted") } diff --git a/src/croc/sender.go b/src/croc/sender.go index 8cce512..9d334c4 100644 --- a/src/croc/sender.go +++ b/src/croc/sender.go @@ -24,7 +24,7 @@ import ( "github.com/schollz/croc/src/utils" "github.com/schollz/croc/src/zipper" "github.com/schollz/pake" - "github.com/schollz/progressbar/v2" + "github.com/schollz/progressbar" "github.com/schollz/spinner" "github.com/tscholl2/siec" ) @@ -46,7 +46,6 @@ func (cr *Croc) startSender(forceSend int, serverAddress string, tcpPorts []stri func (cr *Croc) send(forceSend int, serverAddress string, tcpPorts []string, isLocal bool, c *websocket.Conn, fname string, codephrase string, useCompression bool, useEncryption bool) (err error) { var f *os.File defer f.Close() // ignore the error if it wasn't opened :( - var fstats models.FileStats var fileHash []byte var otherIP string var startTransfer time.Time @@ -137,7 +136,7 @@ func (cr *Croc) send(forceSend int, serverAddress string, tcpPorts []string, isL fileReady <- err return } - fstats = models.FileStats{ + cr.FileInfo = models.FileStats{ Name: filename, Size: fstat.Size(), ModTime: fstat.ModTime(), @@ -146,15 +145,15 @@ func (cr *Croc) send(forceSend int, serverAddress string, tcpPorts []string, isL IsCompressed: useCompression, IsEncrypted: useEncryption, } - if fstats.IsDir { + if cr.FileInfo.IsDir { // zip the directory - fstats.SentName, err = zipper.ZipFile(fname, true) + cr.FileInfo.SentName, err = zipper.ZipFile(fname, true) if err != nil { log.Error(err) fileReady <- err return } - fname = fstats.SentName + fname = cr.FileInfo.SentName fstat, err := os.Stat(fname) if err != nil { @@ -162,7 +161,7 @@ func (cr *Croc) send(forceSend int, serverAddress string, tcpPorts []string, isL return } // get new size - fstats.Size = fstat.Size() + cr.FileInfo.Size = fstat.Size() } // open the file @@ -180,7 +179,7 @@ func (cr *Croc) send(forceSend int, serverAddress string, tcpPorts []string, isL c.WriteMessage(websocket.BinaryMessage, P.Bytes()) // start PAKE spinnner spin.Suffix = " performing PAKE..." - cr.State = "Performing PAKE..." + cr.StateString = "Performing PAKE..." spin.Start() case 2: // P recieves H(k),v from Q @@ -196,7 +195,7 @@ func (cr *Croc) send(forceSend int, serverAddress string, tcpPorts []string, isL // wait for readiness spin.Stop() spin.Suffix = " waiting for recipient ok..." - cr.State = "Waiting for recipient ok...." + cr.StateString = "Waiting for recipient ok...." spin.Start() case 3: log.Debugf("[%d] recipient declares readiness for file info", step) @@ -228,7 +227,7 @@ func (cr *Croc) send(forceSend int, serverAddress string, tcpPorts []string, isL if err != nil { return err } - fstatsBytes, err := json.Marshal(fstats) + fstatsBytes, err := json.Marshal(cr.FileInfo) if err != nil { return err } @@ -267,9 +266,9 @@ func (cr *Croc) send(forceSend int, serverAddress string, tcpPorts []string, isL // start loading the file into memory // start streaming encryption/compression - if fstats.IsDir { + if cr.FileInfo.IsDir { // remove file if zipped - defer os.Remove(fstats.SentName) + defer os.Remove(cr.FileInfo.SentName) } go func(dataChan chan DataChan) { var buffer []byte @@ -291,7 +290,7 @@ func (cr *Croc) send(forceSend int, serverAddress string, tcpPorts []string, isL // do compression var compressedBytes []byte - if useCompression && !fstats.IsDir { + if useCompression && !cr.FileInfo.IsDir { compressedBytes = compress.Compress(buffer[:bytesread]) } else { compressedBytes = buffer[:bytesread] @@ -366,13 +365,13 @@ func (cr *Croc) send(forceSend int, serverAddress string, tcpPorts []string, isL } else { blockSize = models.TCP_BUFFER_SIZE / 2 } - bar := progressbar.NewOptions( - int(fstats.Size), + cr.Bar = progressbar.NewOptions( + int(cr.FileInfo.Size), progressbar.OptionSetRenderBlankState(true), - progressbar.OptionSetBytes(int(fstats.Size)), + progressbar.OptionSetBytes(int(cr.FileInfo.Size)), progressbar.OptionSetWriter(os.Stderr), ) - bar.Add(blockSize * len(blocksToSkip)) + cr.Bar.Add(blockSize * len(blocksToSkip)) if useWebsockets { for { @@ -380,7 +379,7 @@ func (cr *Croc) send(forceSend int, serverAddress string, tcpPorts []string, isL if data.err != nil { return data.err } - bar.Add(data.bytesRead) + cr.Bar.Add(data.bytesRead) // write data to websockets err = c.WriteMessage(websocket.BinaryMessage, data.b) @@ -409,7 +408,7 @@ func (cr *Croc) send(forceSend int, serverAddress string, tcpPorts []string, isL log.Error(data.err) return } - bar.Add(data.bytesRead) + cr.Bar.Add(data.bytesRead) // write data to tcp connection _, err = tcpConnections[i].Write(data.b) if err != nil { @@ -427,7 +426,7 @@ func (cr *Croc) send(forceSend int, serverAddress string, tcpPorts []string, isL wg.Wait() } - bar.Finish() + cr.Bar.Finish() log.Debug("send hash to finish file") fileHash, err = utils.HashFile(fname) if err != nil { @@ -445,10 +444,10 @@ func (cr *Croc) send(forceSend int, serverAddress string, tcpPorts []string, isL log.Debugf("[%d] determing whether it went ok", step) if bytes.Equal(message, fileHash) { log.Debug("file transfered successfully") - transferRate := float64(fstats.Size) / 1000000.0 / transferTime.Seconds() + transferRate := float64(cr.FileInfo.Size) / 1000000.0 / transferTime.Seconds() transferType := "MB/s" if transferRate < 1 { - transferRate = float64(fstats.Size) / 1000.0 / transferTime.Seconds() + transferRate = float64(cr.FileInfo.Size) / 1000.0 / transferTime.Seconds() transferType = "kB/s" } fmt.Fprintf(os.Stderr, "\nTransfer complete (%2.1f %s)", transferRate, transferType)