add file stats to state

This commit is contained in:
Zack Scholl 2018-10-17 06:19:01 -07:00
parent 55d9137b6d
commit 47ab799c6f
3 changed files with 64 additions and 61 deletions

View File

@ -5,8 +5,10 @@ import (
"time" "time"
"github.com/schollz/croc/src/logger" "github.com/schollz/croc/src/logger"
"github.com/schollz/croc/src/models"
"github.com/schollz/croc/src/relay" "github.com/schollz/croc/src/relay"
"github.com/schollz/croc/src/zipper" "github.com/schollz/croc/src/zipper"
"github.com/schollz/progressbar"
) )
func init() { func init() {
@ -51,7 +53,10 @@ type Croc struct {
isLocal bool isLocal bool
normalFinish bool normalFinish bool
State string // state variables
StateString string
Bar *progressbar.ProgressBar
FileInfo models.FileStats
} }
// Init will initiate with the default parameters // Init will initiate with the default parameters

View File

@ -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) { 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 sessionKey []byte
var transferTime time.Duration var transferTime time.Duration
var hash256 []byte var hash256 []byte
@ -172,16 +171,16 @@ func (cr *Croc) receive(forceSend int, serverAddress string, tcpPorts []string,
if err != nil { if err != nil {
return err return err
} }
err = json.Unmarshal(decryptedFileData, &fstats) err = json.Unmarshal(decryptedFileData, &cr.FileInfo)
if err != nil { if err != nil {
return err 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 // 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" overwritingOrReceiving := "Receiving"
if utils.Exists(fstats.Name) || utils.Exists(fstats.SentName) { if utils.Exists(cr.FileInfo.Name) || utils.Exists(cr.FileInfo.SentName) {
overwritingOrReceiving = "Overwriting" overwritingOrReceiving = "Overwriting"
if utils.Exists(progressFile) { if utils.Exists(progressFile) {
overwritingOrReceiving = "Resume receiving" overwritingOrReceiving = "Resume receiving"
@ -205,14 +204,14 @@ func (cr *Croc) receive(forceSend int, serverAddress string, tcpPorts []string,
// prompt user about the file // prompt user about the file
fileOrFolder := "file" fileOrFolder := "file"
if fstats.IsDir { if cr.FileInfo.IsDir {
fileOrFolder = "folder" fileOrFolder = "folder"
} }
fmt.Fprintf(os.Stderr, "\r%s %s (%s) into: %s\n", fmt.Fprintf(os.Stderr, "\r%s %s (%s) into: %s\n",
overwritingOrReceiving, overwritingOrReceiving,
fileOrFolder, fileOrFolder,
humanize.Bytes(uint64(fstats.Size)), humanize.Bytes(uint64(cr.FileInfo.Size)),
fstats.Name, cr.FileInfo.Name,
) )
if !noPrompt { if !noPrompt {
if "y" != utils.GetInput("ok? (y/N): ") { if "y" != utils.GetInput("ok? (y/N): ") {
@ -225,27 +224,27 @@ func (cr *Croc) receive(forceSend int, serverAddress string, tcpPorts []string,
// await file // await file
// erase file if overwriting // erase file if overwriting
if overwritingOrReceiving == "Overwriting" { if overwritingOrReceiving == "Overwriting" {
os.Remove(fstats.SentName) os.Remove(cr.FileInfo.SentName)
} }
var f *os.File var f *os.File
if utils.Exists(fstats.SentName) && resumeFile { if utils.Exists(cr.FileInfo.SentName) && resumeFile {
if !useWebsockets { if !useWebsockets {
f, err = os.OpenFile(fstats.SentName, os.O_WRONLY, 0644) f, err = os.OpenFile(cr.FileInfo.SentName, os.O_WRONLY, 0644)
} else { } 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 { if err != nil {
log.Error(err) log.Error(err)
return err return err
} }
} else { } else {
f, err = os.Create(fstats.SentName) f, err = os.Create(cr.FileInfo.SentName)
if err != nil { if err != nil {
log.Error(err) log.Error(err)
return err return err
} }
if !useWebsockets { if !useWebsockets {
if err = f.Truncate(fstats.Size); err != nil { if err = f.Truncate(cr.FileInfo.Size); err != nil {
log.Error(err) log.Error(err)
return err return err
} }
@ -262,13 +261,13 @@ func (cr *Croc) receive(forceSend int, serverAddress string, tcpPorts []string,
// start the ui for pgoress // start the ui for pgoress
bytesWritten := 0 bytesWritten := 0
fmt.Fprintf(os.Stderr, "\nReceiving (<-%s)...\n", otherIP) fmt.Fprintf(os.Stderr, "\nReceiving (<-%s)...\n", otherIP)
bar := progressbar.NewOptions( cr.Bar = progressbar.NewOptions(
int(fstats.Size), int(cr.FileInfo.Size),
progressbar.OptionSetRenderBlankState(true), progressbar.OptionSetRenderBlankState(true),
progressbar.OptionSetBytes(int(fstats.Size)), progressbar.OptionSetBytes(int(cr.FileInfo.Size)),
progressbar.OptionSetWriter(os.Stderr), progressbar.OptionSetWriter(os.Stderr),
) )
bar.Add((len(blocks) * blockSize)) cr.Bar.Add((len(blocks) * blockSize))
finished := make(chan bool) finished := make(chan bool)
go func(finished chan bool, dataChan chan []byte) (err error) { 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() defer fProgress.Close()
blocksWritten := 0.0 blocksWritten := 0.0
blocksToWrite := float64(fstats.Size) blocksToWrite := float64(cr.FileInfo.Size)
if useWebsockets { if useWebsockets {
blocksToWrite = blocksToWrite/float64(models.WEBSOCKET_BUFFER_SIZE/8) - float64(len(blocks)) blocksToWrite = blocksToWrite/float64(models.WEBSOCKET_BUFFER_SIZE/8) - float64(len(blocks))
} else { } else {
@ -304,7 +303,7 @@ func (cr *Croc) receive(forceSend int, serverAddress string, tcpPorts []string,
log.Error(err) log.Error(err)
return err return err
} }
decrypted, err := enc.Decrypt(sessionKey, !fstats.IsEncrypted) decrypted, err := enc.Decrypt(sessionKey, !cr.FileInfo.IsEncrypted)
if err != nil { if err != nil {
log.Error(err) log.Error(err)
return err return err
@ -319,7 +318,7 @@ func (cr *Croc) receive(forceSend int, serverAddress string, tcpPorts []string,
} }
// do decompression // do decompression
if fstats.IsCompressed && !fstats.IsDir { if cr.FileInfo.IsCompressed && !cr.FileInfo.IsDir {
decrypted = compress.Decompress(decrypted) decrypted = compress.Decompress(decrypted)
} }
@ -347,9 +346,9 @@ func (cr *Croc) receive(forceSend int, serverAddress string, tcpPorts []string,
bytesWritten += n bytesWritten += n
blocksWritten += 1.0 blocksWritten += 1.0
// update the progress bar // update the progress bar
bar.Add(n) cr.Bar.Add(n)
if int64(bytesWritten) == fstats.Size || blocksWritten >= blocksToWrite { if int64(bytesWritten) == cr.FileInfo.Size || blocksWritten >= blocksToWrite {
log.Debug("finished", int64(bytesWritten), fstats.Size, blocksWritten, blocksToWrite) log.Debug("finished", int64(bytesWritten), cr.FileInfo.Size, blocksWritten, blocksToWrite)
break break
} }
} }
@ -432,10 +431,10 @@ func (cr *Croc) receive(forceSend int, serverAddress string, tcpPorts []string,
} }
// finish bar // finish bar
bar.Finish() cr.Bar.Finish()
// check hash // check hash
hash256, err = utils.HashFile(fstats.SentName) hash256, err = utils.HashFile(cr.FileInfo.SentName)
if err != nil { if err != nil {
log.Error(err) log.Error(err)
return err return err
@ -447,45 +446,45 @@ func (cr *Croc) receive(forceSend int, serverAddress string, tcpPorts []string,
log.Debugf("got hash: %x", message) log.Debugf("got hash: %x", message)
if bytes.Equal(hash256, message) { if bytes.Equal(hash256, message) {
// open directory // open directory
if fstats.IsDir { if cr.FileInfo.IsDir {
err = zipper.UnzipFile(fstats.SentName, ".") err = zipper.UnzipFile(cr.FileInfo.SentName, ".")
if DebugLevel != "debug" { if DebugLevel != "debug" {
os.Remove(fstats.SentName) os.Remove(cr.FileInfo.SentName)
} }
} else { } else {
err = nil err = nil
} }
if err == nil { if err == nil {
if useStdout && !fstats.IsDir { if useStdout && !cr.FileInfo.IsDir {
var bFile []byte var bFile []byte
bFile, err = ioutil.ReadFile(fstats.SentName) bFile, err = ioutil.ReadFile(cr.FileInfo.SentName)
if err != nil { if err != nil {
return err return err
} }
os.Stdout.Write(bFile) 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" transferType := "MB/s"
if transferRate < 1 { if transferRate < 1 {
transferRate = float64(fstats.Size) / 1000.0 / transferTime.Seconds() transferRate = float64(cr.FileInfo.Size) / 1000.0 / transferTime.Seconds()
transferType = "kB/s" transferType = "kB/s"
} }
folderOrFile := "file" folderOrFile := "file"
if fstats.IsDir { if cr.FileInfo.IsDir {
folderOrFile = "folder" folderOrFile = "folder"
} }
if useStdout { 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) os.Remove(progressFile)
} }
return err return err
} else { } else {
if DebugLevel != "debug" { if DebugLevel != "debug" {
log.Debug("removing corrupted file") log.Debug("removing corrupted file")
os.Remove(fstats.SentName) os.Remove(cr.FileInfo.SentName)
} }
return errors.New("file corrupted") return errors.New("file corrupted")
} }

View File

@ -24,7 +24,7 @@ import (
"github.com/schollz/croc/src/utils" "github.com/schollz/croc/src/utils"
"github.com/schollz/croc/src/zipper" "github.com/schollz/croc/src/zipper"
"github.com/schollz/pake" "github.com/schollz/pake"
"github.com/schollz/progressbar/v2" "github.com/schollz/progressbar"
"github.com/schollz/spinner" "github.com/schollz/spinner"
"github.com/tscholl2/siec" "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) { 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 var f *os.File
defer f.Close() // ignore the error if it wasn't opened :( defer f.Close() // ignore the error if it wasn't opened :(
var fstats models.FileStats
var fileHash []byte var fileHash []byte
var otherIP string var otherIP string
var startTransfer time.Time var startTransfer time.Time
@ -137,7 +136,7 @@ func (cr *Croc) send(forceSend int, serverAddress string, tcpPorts []string, isL
fileReady <- err fileReady <- err
return return
} }
fstats = models.FileStats{ cr.FileInfo = models.FileStats{
Name: filename, Name: filename,
Size: fstat.Size(), Size: fstat.Size(),
ModTime: fstat.ModTime(), ModTime: fstat.ModTime(),
@ -146,15 +145,15 @@ func (cr *Croc) send(forceSend int, serverAddress string, tcpPorts []string, isL
IsCompressed: useCompression, IsCompressed: useCompression,
IsEncrypted: useEncryption, IsEncrypted: useEncryption,
} }
if fstats.IsDir { if cr.FileInfo.IsDir {
// zip the directory // zip the directory
fstats.SentName, err = zipper.ZipFile(fname, true) cr.FileInfo.SentName, err = zipper.ZipFile(fname, true)
if err != nil { if err != nil {
log.Error(err) log.Error(err)
fileReady <- err fileReady <- err
return return
} }
fname = fstats.SentName fname = cr.FileInfo.SentName
fstat, err := os.Stat(fname) fstat, err := os.Stat(fname)
if err != nil { if err != nil {
@ -162,7 +161,7 @@ func (cr *Croc) send(forceSend int, serverAddress string, tcpPorts []string, isL
return return
} }
// get new size // get new size
fstats.Size = fstat.Size() cr.FileInfo.Size = fstat.Size()
} }
// open the file // open the file
@ -180,7 +179,7 @@ func (cr *Croc) send(forceSend int, serverAddress string, tcpPorts []string, isL
c.WriteMessage(websocket.BinaryMessage, P.Bytes()) c.WriteMessage(websocket.BinaryMessage, P.Bytes())
// start PAKE spinnner // start PAKE spinnner
spin.Suffix = " performing PAKE..." spin.Suffix = " performing PAKE..."
cr.State = "Performing PAKE..." cr.StateString = "Performing PAKE..."
spin.Start() spin.Start()
case 2: case 2:
// P recieves H(k),v from Q // 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 // wait for readiness
spin.Stop() spin.Stop()
spin.Suffix = " waiting for recipient ok..." spin.Suffix = " waiting for recipient ok..."
cr.State = "Waiting for recipient ok...." cr.StateString = "Waiting for recipient ok...."
spin.Start() spin.Start()
case 3: case 3:
log.Debugf("[%d] recipient declares readiness for file info", step) 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 { if err != nil {
return err return err
} }
fstatsBytes, err := json.Marshal(fstats) fstatsBytes, err := json.Marshal(cr.FileInfo)
if err != nil { if err != nil {
return err return err
} }
@ -267,9 +266,9 @@ func (cr *Croc) send(forceSend int, serverAddress string, tcpPorts []string, isL
// start loading the file into memory // start loading the file into memory
// start streaming encryption/compression // start streaming encryption/compression
if fstats.IsDir { if cr.FileInfo.IsDir {
// remove file if zipped // remove file if zipped
defer os.Remove(fstats.SentName) defer os.Remove(cr.FileInfo.SentName)
} }
go func(dataChan chan DataChan) { go func(dataChan chan DataChan) {
var buffer []byte var buffer []byte
@ -291,7 +290,7 @@ func (cr *Croc) send(forceSend int, serverAddress string, tcpPorts []string, isL
// do compression // do compression
var compressedBytes []byte var compressedBytes []byte
if useCompression && !fstats.IsDir { if useCompression && !cr.FileInfo.IsDir {
compressedBytes = compress.Compress(buffer[:bytesread]) compressedBytes = compress.Compress(buffer[:bytesread])
} else { } else {
compressedBytes = buffer[:bytesread] compressedBytes = buffer[:bytesread]
@ -366,13 +365,13 @@ func (cr *Croc) send(forceSend int, serverAddress string, tcpPorts []string, isL
} else { } else {
blockSize = models.TCP_BUFFER_SIZE / 2 blockSize = models.TCP_BUFFER_SIZE / 2
} }
bar := progressbar.NewOptions( cr.Bar = progressbar.NewOptions(
int(fstats.Size), int(cr.FileInfo.Size),
progressbar.OptionSetRenderBlankState(true), progressbar.OptionSetRenderBlankState(true),
progressbar.OptionSetBytes(int(fstats.Size)), progressbar.OptionSetBytes(int(cr.FileInfo.Size)),
progressbar.OptionSetWriter(os.Stderr), progressbar.OptionSetWriter(os.Stderr),
) )
bar.Add(blockSize * len(blocksToSkip)) cr.Bar.Add(blockSize * len(blocksToSkip))
if useWebsockets { if useWebsockets {
for { for {
@ -380,7 +379,7 @@ func (cr *Croc) send(forceSend int, serverAddress string, tcpPorts []string, isL
if data.err != nil { if data.err != nil {
return data.err return data.err
} }
bar.Add(data.bytesRead) cr.Bar.Add(data.bytesRead)
// write data to websockets // write data to websockets
err = c.WriteMessage(websocket.BinaryMessage, data.b) 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) log.Error(data.err)
return return
} }
bar.Add(data.bytesRead) cr.Bar.Add(data.bytesRead)
// write data to tcp connection // write data to tcp connection
_, err = tcpConnections[i].Write(data.b) _, err = tcpConnections[i].Write(data.b)
if err != nil { if err != nil {
@ -427,7 +426,7 @@ func (cr *Croc) send(forceSend int, serverAddress string, tcpPorts []string, isL
wg.Wait() wg.Wait()
} }
bar.Finish() cr.Bar.Finish()
log.Debug("send hash to finish file") log.Debug("send hash to finish file")
fileHash, err = utils.HashFile(fname) fileHash, err = utils.HashFile(fname)
if err != nil { 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) log.Debugf("[%d] determing whether it went ok", step)
if bytes.Equal(message, fileHash) { if bytes.Equal(message, fileHash) {
log.Debug("file transfered successfully") 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" transferType := "MB/s"
if transferRate < 1 { if transferRate < 1 {
transferRate = float64(fstats.Size) / 1000.0 / transferTime.Seconds() transferRate = float64(cr.FileInfo.Size) / 1000.0 / transferTime.Seconds()
transferType = "kB/s" transferType = "kB/s"
} }
fmt.Fprintf(os.Stderr, "\nTransfer complete (%2.1f %s)", transferRate, transferType) fmt.Fprintf(os.Stderr, "\nTransfer complete (%2.1f %s)", transferRate, transferType)