go-examples/zip/zip.go

package main

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

// Zip walks `src`, omitting `trim`, writing to `w`
func Zip(w io.Writer, src string, trim string) error {
	zw := zip.NewWriter(w)
	defer zw.Close()

	return filepath.Walk(src, func(path string, fi os.FileInfo, err error) error {
		// path includes fi.Name() already
		if nil != err {
			fmt.Println("warning: skipped", path+": ", err)
			return nil
		}

		zipOne(zw, path, fi, trim)
		return nil
	})
}

func zipOne(zw *zip.Writer, path string, fi os.FileInfo, trim string) error {
	h, err := zip.FileInfoHeader(fi)
	if nil != err {
		return err
	}
	h.Name = strings.TrimPrefix(strings.TrimPrefix(path, trim), string(filepath.Separator))

	if fi.IsDir() {
		fmt.Printf("directory: %s\n\t%q\n", path, h.Name)
		return zipDirectory(zw, h)
	}

	// Allow zipping a single file
	if "" == h.Name {
		h.Name = path
	}
	if fi.Mode().IsRegular() {
		fmt.Printf("file: %s\n\t%q\n", path, h.Name)
		return zipFile(zw, h, path)
	}

	if os.ModeSymlink == (fi.Mode() & os.ModeType) {
		fmt.Printf("symlink: %s\n\t%q\n", path, h.Name)
		return zipSymlink(zw, h, path)
	}

	fmt.Printf("skipping: %s\n\t(irregular file type)\n", path)
	return nil
}

func zipDirectory(zw *zip.Writer, h *zip.FileHeader) error {
	// directories must end in / for go
	h.Name = strings.TrimPrefix(h.Name+"/", "/")

	// skip top-level, trimmed directory
	if "" == h.Name {
		return nil
	}

	if _, err := zw.CreateHeader(h); nil != err {
		return err
	}

	return nil
}

func zipFile(zw *zip.Writer, h *zip.FileHeader, path string) error {
	r, err := os.Open(path)
	if nil != err {
		return err
	}
	defer r.Close()

	// Files should be zipped (not dirs, and symlinks... meh)
	// TODO investigate if files below a certain size shouldn't be deflated
	h.Method = zip.Deflate
	w, err := zw.CreateHeader(h)
	if nil != err {
		return err
	}

	if _, err := io.Copy(w, r); nil != err {
		return err
	}

	return nil
}

func zipSymlink(zw *zip.Writer, h *zip.FileHeader, path string) error {
	w, err := zw.CreateHeader(h)
	if nil != err {
		return err
	}

	// TODO make sure that this is within the root directory
	targetpath, err := os.Readlink(path)
	if nil != err {
		return err
	}
	if _, err := w.Write([]byte(targetpath)); nil != err {
		return err
	}

	return nil
}