bitbucket.org/ai69/amoy@v0.2.3/zip.go

package amoy

import (
	"archive/zip"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"time"
)

// ArchiveContent represents a map between filename and data.
type ArchiveContent map[string][]byte

//revive:disable:error-naming It's not a real error
var (
	// QuitUnzip indicates the arbitrary error means to quit from unzip.
	QuitUnzip = errors.New("amoy: quit unzip file")
)

// ZipDir compresses one or many directories into a single zip archive file. Existing destination file will be overwritten.
func ZipDir(destZip string, srcDirs ...string) error {
	fw, err := os.Create(destZip)
	if err != nil {
		return err
	}
	defer fw.Close()

	zw := zip.NewWriter(fw)
	defer zw.Close()

	for _, d := range srcDirs {
		if err := addDirToZip(zw, d); err != nil {
			return err
		}
	}
	return nil
}

// ZipFile compresses one or many files into a single zip archive file.
func ZipFile(destZip string, srcFiles ...string) error {
	fw, err := os.Create(destZip)
	if err != nil {
		return err
	}
	defer fw.Close()

	zw := zip.NewWriter(fw)
	defer zw.Close()

	for _, f := range srcFiles {
		if err := addFileToZip(zw, f); err != nil {
			return err
		}
	}
	return nil
}

// ZipContent compresses data entries into a single zip archive file.
func ZipContent(destZip string, content ArchiveContent) error {
	fw, err := os.Create(destZip)
	if err != nil {
		return err
	}
	defer fw.Close()

	zw := zip.NewWriter(fw)
	defer zw.Close()

	for name, data := range content {
		if err := addContentToZip(zw, name, data); err != nil {
			return err
		}
	}
	return nil
}

func addDirToZip(zw *zip.Writer, dirname string) error {
	baseDir := filepath.Base(dirname)
	err := filepath.Walk(dirname, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		header, err := zip.FileInfoHeader(info)
		if err != nil {
			return err
		}

		// get relative entry name
		switch baseDir {
		case ".":
			header.Name = path
		case "..":
			header.Name = strings.TrimLeft(strings.TrimPrefix(path, dirname), `\/`)
		default:
			header.Name = filepath.Join(baseDir, strings.TrimPrefix(path, dirname))
		}

		// skip directory like ".", "..", "../.."
		if header.Name == "" || header.Name == "." {
			return nil
		}

		header.Name = filepath.ToSlash(header.Name)
		if info.IsDir() {
			header.Name += "/"
			// create directory entry
			if _, err := zw.CreateHeader(header); err != nil {
				return err
			}
			return nil
		}

		header.Method = zip.Deflate
		// open source file
		fr, err := os.Open(path)
		if err != nil {
			return err
		}
		defer fr.Close()

		// create file entry
		fw, err := zw.CreateHeader(header)
		if err != nil {
			return err
		}
		_, err = io.Copy(fw, fr)
		return err
	})
	return err
}

func addFileToZip(zw *zip.Writer, filename string) error {
	fr, err := os.Open(filename)
	if err != nil {
		return err
	}
	defer fr.Close()

	info, err := fr.Stat()
	if err != nil {
		return err
	}

	header, err := zip.FileInfoHeader(info)
	if err != nil {
		return err
	}
	header.Name = filepath.ToSlash(filename)
	header.Method = zip.Deflate
	fw, err := zw.CreateHeader(header)
	if err != nil {
		return err
	}

	_, err = io.Copy(fw, fr)
	return err
}

func addContentToZip(zw *zip.Writer, filename string, data []byte) error {
	header := new(zip.FileHeader)
	header.Name = filepath.ToSlash(filename)
	header.Method = zip.Deflate
	header.Modified = time.Now()
	header.UncompressedSize64 = uint64(len(data))
	if header.UncompressedSize64 > uint64(MaxUint32) {
		header.UncompressedSize = MaxUint32
	} else {
		header.UncompressedSize = uint32(header.UncompressedSize64)
	}

	fw, err := zw.CreateHeader(header)
	if err != nil {
		return err
	}
	_, err = fw.Write(data)
	return err
}

// UnzipConflictStrategy defines the strategy to handle the conflict when extracting files from a zip archive.
type UnzipConflictStrategy uint8

const (
	// UnzipConflictSkip skips the file if it already exists in the destination folder.
	UnzipConflictSkip UnzipConflictStrategy = iota
	// UnzipConflictOverwrite overwrites the file if it already exists in the destination folder.
	UnzipConflictOverwrite
	// UnzipConflictRename renames the new file if it already exists in the destination folder.
	UnzipConflictRename
	// UnzipConflictKeepOld keeps the file with earlier modification time if it already exists in the destination folder.
	UnzipConflictKeepOld
	// UnzipConflictKeepNew keeps the file with later modification time if it already exists in the destination folder.
	UnzipConflictKeepNew
	maxCountUnzipConflict
)

// UnzipDir decompresses a zip archive, extracts all files and folders within the zip file to an output directory.
func UnzipDir(srcZip, destDir string, opts ...UnzipConflictStrategy) ([]string, error) {
	var filenames []string

	// Conflict strategy
	strategy := UnzipConflictSkip
	if len(opts) > 0 {
		for _, o := range opts {
			if o < maxCountUnzipConflict {
				strategy = o
				break
			}
		}
	}

	zr, err := zip.OpenReader(srcZip)
	if err != nil {
		return filenames, err
	}
	defer zr.Close()

	// explicit the destination folder
	destDir, _ = filepath.Abs(destDir)

	for _, f := range zr.File {
		// Store filename/path for returning and using later on
		fpath := filepath.Join(destDir, f.Name)
		mt := f.Modified

		// Check for ZipSlip. More Info: http://bit.ly/2MsjAWE
		if !strings.HasPrefix(fpath, filepath.Clean(destDir)+string(os.PathSeparator)) {
			return filenames, fmt.Errorf("%s: illegal file path", fpath)
		}

		// Create directory and skip
		if f.FileInfo().IsDir() {
			// Make Folder
			if err := os.MkdirAll(fpath, f.Mode()); err != nil {
				return filenames, err
			}
			continue
		}

		// Make parent directory if not exist
		filenames = append(filenames, fpath)
		if err = os.MkdirAll(filepath.Dir(fpath), os.ModePerm); err != nil {
			return filenames, err
		}

		// Check if path already exists
		fi, err := os.Stat(fpath)
		if err == nil {
			// Skip if it is a directory
			if fi.IsDir() {
				continue
			}
			// If file already exists, check the conflict strategy
			switch strategy {
			case UnzipConflictSkip:
				continue
			case UnzipConflictOverwrite:
			case UnzipConflictRename:
				fpath = renameConflictPath(fpath)
			case UnzipConflictKeepOld:
				if fi.ModTime().Before(mt) {
					continue
				}
			case UnzipConflictKeepNew:
				if fi.ModTime().After(mt) {
					continue
				}
			}
		} else if !os.IsNotExist(err) {
			// If file does not exist, continue, returns error if it is other error
			return filenames, err
		}

		// Create new file for writing
		fw, err := os.OpenFile(fpath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, f.Mode())
		if err != nil {
			return filenames, err
		}

		// Open zip entry for reading
		fr, err := f.Open()
		if err != nil {
			return filenames, err
		}

		// Write files
		_, err = io.Copy(fw, fr)

		// Close the file without defer to close before next iteration of loop
		_ = fw.Close()
		_ = fr.Close()

		// Set modification time, timezone offset may be lost for zip file created on Windows
		_ = os.Chtimes(fpath, mt, mt)

		if err != nil {
			return filenames, err
		}
	}

	return filenames, nil
}

// renameConflictPath renames the path if it already exists.
func renameConflictPath(path string) string {
	// check if the path exists
	_, err := os.Stat(path)
	if err != nil && os.IsNotExist(err) {
		// if it does not exist, return the original path
		return path
	}

	// for other cases, add a suffix to the path and check again
	ext := filepath.Ext(path)
	name := strings.TrimSuffix(path, ext)
	for i := 1; ; i++ {
		newPath := fmt.Sprintf("%s-%d%s", name, i, ext)
		_, err := os.Stat(newPath)
		if err == nil {
			// skip if it exists, no matter it is a file or a directory
			continue
		}
		if os.IsNotExist(err) {
			// great! the path does not exist
			return newPath
		}
	}
}

// UnzipFile decompresses a zip archive, extracts all files and call handlers.
func UnzipFile(srcZip string, handle func(file *zip.File) error) error {
	zr, err := zip.OpenReader(srcZip)
	if err != nil {
		return err
	}
	defer zr.Close()

	for _, f := range zr.File {
		// skip directory
		if f.FileInfo().IsDir() {
			continue
		}

		// call handler
		if err = handle(f); err != nil {
			break
		}
	}

	if err == QuitUnzip {
		err = nil
	}
	return err
}

// UnzipContent decompresses a zip archive, extracts all files within the zip file to map of bytes.
func UnzipContent(srcZip string) (ArchiveContent, error) {
	store := make(ArchiveContent)
	err := UnzipFile(srcZip, func(f *zip.File) error {
		// Open file
		fr, err := f.Open()
		if err != nil {
			return err
		}
		defer fr.Close()

		// Read content
		bytes, err := ioutil.ReadAll(fr)
		if err == nil {
			store[f.Name] = bytes
		}
		return err
	})
	return store, err
}