From caf551180534b7fdf1f6fc77b6da1bfa230cda04 Mon Sep 17 00:00:00 2001 From: Zack Scholl Date: Fri, 21 Sep 2018 21:50:20 -0700 Subject: [PATCH] add zipping --- src/croc/croc.go | 2 + src/sender/sender.go | 27 +++--- src/zipper/zip.go | 196 +++++++++++++++++++++++++++++++++++++++++ src/zipper/zip_test.go | 23 +++++ 4 files changed, 237 insertions(+), 11 deletions(-) create mode 100644 src/zipper/zip.go create mode 100644 src/zipper/zip_test.go diff --git a/src/croc/croc.go b/src/croc/croc.go index 373ded4..7cc236f 100644 --- a/src/croc/croc.go +++ b/src/croc/croc.go @@ -7,6 +7,7 @@ import ( "github.com/schollz/croc/src/recipient" "github.com/schollz/croc/src/relay" "github.com/schollz/croc/src/sender" + "github.com/schollz/croc/src/zipper" ) // Croc options @@ -61,5 +62,6 @@ func Init(debug bool) (c *Croc) { sender.DebugLevel = debugLevel recipient.DebugLevel = debugLevel relay.DebugLevel = debugLevel + zipper.DebugLevel = debugLevel return } diff --git a/src/sender/sender.go b/src/sender/sender.go index dac242a..f956163 100644 --- a/src/sender/sender.go +++ b/src/sender/sender.go @@ -39,8 +39,22 @@ func Send(done chan struct{}, c *websocket.Conn, fname string, codephrase string } func send(c *websocket.Conn, fname string, codephrase string) (err error) { - var f *os.File - var fstats models.FileStats + // check that the file exists + f, err := os.Open(fname) + if err != nil { + return + } + fstat, err := f.Stat() + if err != nil { + return err + } + fstats := models.FileStats{fstat.Name(), fstat.Size(), fstat.ModTime(), fstat.IsDir()} + if fstats.IsDir { + // zip the directory + + } + + // get ready to generate session key var sessionKey []byte // pick an elliptic curve @@ -82,15 +96,6 @@ func send(c *websocket.Conn, fname string, codephrase string) (err error) { if !bytes.Equal(message, []byte("ready")) { return errors.New("recipient refused file") } - f, err = os.Open(fname) - if err != nil { - return - } - fstat, err := f.Stat() - if err != nil { - return err - } - fstats = models.FileStats{fstat.Name(), fstat.Size(), fstat.ModTime(), fstat.IsDir()} fstatsBytes, err := json.Marshal(fstats) if err != nil { return err diff --git a/src/zipper/zip.go b/src/zipper/zip.go new file mode 100644 index 0000000..aff6e29 --- /dev/null +++ b/src/zipper/zip.go @@ -0,0 +1,196 @@ +package zipper + +import ( + "archive/zip" + "compress/flate" + "io" + "io/ioutil" + "os" + "path/filepath" + "strings" + + log "github.com/cihub/seelog" + "github.com/schollz/croc/src/logger" +) + +var DebugLevel string + +func init() { + DebugLevel = "info" +} + +// UnzipFile will unzip the src directory into the dest +func UnzipFile(src, dest string) (err error) { + logger.SetLogLevel(DebugLevel) + + r, err := zip.OpenReader(src) + if err != nil { + return + } + defer r.Close() + + for _, f := range r.File { + var rc io.ReadCloser + rc, err = f.Open() + if err != nil { + return + } + defer rc.Close() + + // Store filename/path for returning and using later on + fpath := filepath.Join(dest, f.Name) + log.Debugf("unzipping %s", fpath) + fpath = filepath.FromSlash(fpath) + + if f.FileInfo().IsDir() { + + // Make Folder + os.MkdirAll(fpath, os.ModePerm) + + } else { + + // Make File + if err = os.MkdirAll(filepath.Dir(fpath), os.ModePerm); err != nil { + return + } + + var outFile *os.File + outFile, err = os.OpenFile(fpath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, f.Mode()) + if err != nil { + return + } + + _, err = io.Copy(outFile, rc) + + // Close the file without defer to close before next iteration of loop + outFile.Close() + + if err != nil { + return + } + + } + } + if err == nil { + log.Debugf("unzipped %s to %s", src, dest) + } + return +} + +// ZipFile will zip the file +func ZipFile(fname string, compress bool) (writtenFilename string, err error) { + logger.SetLogLevel(DebugLevel) + + log.Debugf("zipping %s with compression? %v", fname, compress) + pathtofile, filename := filepath.Split(fname) + curdir, err := os.Getwd() + if err != nil { + log.Error(err) + return + } + curdir, err = filepath.Abs(curdir) + if err != nil { + log.Error(err) + return + } + log.Debugf("current directory: %s", curdir) + newfile, err := ioutil.TempFile(curdir, filename+".") + if err != nil { + log.Error(err) + return + } + writtenFilename = newfile.Name() + defer newfile.Close() + + defer os.Chdir(curdir) + log.Debugf("changing dir to %s", pathtofile) + os.Chdir(pathtofile) + + zipWriter := zip.NewWriter(newfile) + zipWriter.RegisterCompressor(zip.Deflate, func(out io.Writer) (io.WriteCloser, error) { + if compress { + return flate.NewWriter(out, flate.BestSpeed) + } else { + return flate.NewWriter(out, flate.NoCompression) + } + }) + defer zipWriter.Close() + + zipfile, err := os.Open(filename) + if err != nil { + log.Error(err) + return "", err + } + defer zipfile.Close() + // Get the file information + info, err := zipfile.Stat() + if err != nil { + log.Error(err) + return + } + + // write header informaiton + header, err := zip.FileInfoHeader(info) + if err != nil { + log.Error(err) + return + } + + var writer io.Writer + if info.IsDir() { + baseDir := filename + filepath.Walk(baseDir, func(path string, info os.FileInfo, err error) error { + if err != nil { + return err + } + + header, err := zip.FileInfoHeader(info) + if err != nil { + return err + } + + if baseDir != "" { + header.Name = filepath.Join(baseDir, strings.TrimPrefix(path, baseDir)) + } + + if info.IsDir() { + header.Name += "/" + } else { + header.Method = zip.Deflate + } + + header.Name = filepath.ToSlash(header.Name) + + writer, err = zipWriter.CreateHeader(header) + if err != nil { + return err + } + + if info.IsDir() { + return nil + } + + file, err := os.Open(path) + if err != nil { + return err + } + defer file.Close() + _, err = io.Copy(writer, file) + return err + }) + } else { + writer, err = zipWriter.CreateHeader(header) + if err != nil { + log.Error(err) + return + } + _, err = io.Copy(writer, zipfile) + if err != nil { + log.Error(err) + return + } + } + + log.Debugf("wrote zip file to %s", writtenFilename) + return +} diff --git a/src/zipper/zip_test.go b/src/zipper/zip_test.go new file mode 100644 index 0000000..2d6edec --- /dev/null +++ b/src/zipper/zip_test.go @@ -0,0 +1,23 @@ +package zipper + +import ( + "os" + "testing" + + log "github.com/cihub/seelog" + "github.com/schollz/croc/src/utils" + "github.com/stretchr/testify/assert" +) + +func TestZip(t *testing.T) { + defer log.Flush() + DebugLevel = "debug" + writtenFilename, err := ZipFile("../../README.md", false) + assert.Nil(t, err) + defer os.Remove(writtenFilename) + + err = UnzipFile(writtenFilename, ".") + assert.Nil(t, err) + assert.True(t, utils.Exists("README.md")) + os.RemoveAll("README.md") +}