github.com/mweagle/Sparta@v1.15.0/zip/zip.go (about)

     1  package zip
     2  
     3  import (
     4  	"archive/zip"
     5  	"fmt"
     6  	"io"
     7  	"os"
     8  	"path/filepath"
     9  	"strings"
    10  
    11  	"github.com/pkg/errors"
    12  	"github.com/sirupsen/logrus"
    13  )
    14  
    15  // FileHeaderAnnotator represents a callback function that accepts the current
    16  // file being added to allow it to customize the ZIP archive values
    17  type FileHeaderAnnotator func(header *zip.FileHeader) (*zip.FileHeader, error)
    18  
    19  // AnnotateAddToZip is an extended Zip writer that accepts an annotation function
    20  // to customize the FileHeader values written into the archive
    21  func AnnotateAddToZip(zipWriter *zip.Writer,
    22  	source string,
    23  	rootSource string,
    24  	annotator FileHeaderAnnotator,
    25  	logger *logrus.Logger) error {
    26  
    27  	linuxZipName := func(platformValue string) string {
    28  		return strings.Replace(platformValue, "\\", "/", -1)
    29  	}
    30  
    31  	fullPathSource, err := filepath.Abs(source)
    32  	if nil != err {
    33  		return errors.Wrapf(err, "Failed to get absolute filepath")
    34  	}
    35  
    36  	appendFile := func(info os.FileInfo) error {
    37  		zipEntryName := info.Name()
    38  		if rootSource != "" {
    39  			zipEntryName = fmt.Sprintf("%s/%s", linuxZipName(rootSource), info.Name())
    40  		}
    41  		// Create a header for this zipFile, basically let's see
    42  		// if we can get the executable bits to travel along..
    43  		fileHeader, fileHeaderErr := zip.FileInfoHeader(info)
    44  		if fileHeaderErr != nil {
    45  			return fileHeaderErr
    46  		}
    47  		// Update the name to the proper thing...
    48  		fileHeader.Name = zipEntryName
    49  		if annotator != nil {
    50  			annotatedHeader, annotatedHeaderErr := annotator(fileHeader)
    51  			if annotatedHeaderErr != nil {
    52  				return errors.Wrapf(annotatedHeaderErr, "Failed to annotate Zip entry file header")
    53  			}
    54  			fileHeader = annotatedHeader
    55  		}
    56  
    57  		logger.WithFields(logrus.Fields{
    58  			"Source": info.Name(),
    59  			"Header": fileHeader,
    60  		}).Debug("Adding ZIP")
    61  
    62  		// File info for the binary executable
    63  		binaryWriter, binaryWriterErr := zipWriter.CreateHeader(fileHeader)
    64  		if binaryWriterErr != nil {
    65  			return binaryWriterErr
    66  		}
    67  		/* #nosec */
    68  		reader, readerErr := os.Open(fullPathSource)
    69  		if readerErr != nil {
    70  			return readerErr
    71  		}
    72  		written, copyErr := io.Copy(binaryWriter, reader)
    73  		errClose := reader.Close()
    74  		if errClose != nil {
    75  			logger.WithFields(logrus.Fields{
    76  				"Error": errClose,
    77  			}).Warn("Failed to close Zip input stream")
    78  		}
    79  		logger.WithFields(logrus.Fields{
    80  			"WrittenBytes": written,
    81  			"SourcePath":   fullPathSource,
    82  			"ZipName":      zipEntryName,
    83  		}).Debug("Archiving file")
    84  		return copyErr
    85  	}
    86  
    87  	directoryWalker := func(path string, info os.FileInfo, err error) error {
    88  		if err != nil {
    89  			return err
    90  		}
    91  		header, err := zip.FileInfoHeader(info)
    92  		if err != nil {
    93  			return errors.Wrapf(err, "Failed to create FileInfoHeader")
    94  		}
    95  		// Normalize the Name
    96  		platformName := strings.TrimPrefix(strings.TrimPrefix(path, rootSource), string(os.PathSeparator))
    97  		header.Name = linuxZipName(platformName)
    98  
    99  		if info.IsDir() {
   100  			header.Name += "/"
   101  		} else {
   102  			header.Method = zip.Deflate
   103  		}
   104  		writer, err := zipWriter.CreateHeader(header)
   105  		if err != nil {
   106  			return err
   107  		}
   108  		if info.IsDir() {
   109  			return nil
   110  		}
   111  		/* #nosec */
   112  		file, err := os.Open(path)
   113  		if err != nil {
   114  			return errors.Wrapf(err, "Failed to open file: %s", path)
   115  		}
   116  		defer func() {
   117  			closeErr := file.Close()
   118  			if closeErr != nil {
   119  				logger.WithFields(logrus.Fields{
   120  					"error": closeErr,
   121  				}).Warn("Failed to close zip writer")
   122  			}
   123  		}()
   124  
   125  		_, err = io.Copy(writer, file)
   126  		if err != nil {
   127  			return errors.Wrapf(err, "Failed to copy file contents")
   128  		}
   129  		return nil
   130  	}
   131  
   132  	fileInfo, err := os.Stat(fullPathSource)
   133  	if nil != err {
   134  		return errors.Wrapf(err, "Failed to get file information")
   135  	}
   136  	switch mode := fileInfo.Mode(); {
   137  	case mode.IsDir():
   138  		err = filepath.Walk(fullPathSource, directoryWalker)
   139  	case mode.IsRegular():
   140  		err = appendFile(fileInfo)
   141  	default:
   142  		err = errors.New("Inavlid source type")
   143  	}
   144  	if err != nil {
   145  		return errors.Wrapf(err, "Failed to determine file mode")
   146  	}
   147  	return nil
   148  }
   149  
   150  // AddToZip creates a source object (either a file, or a directory that will be recursively
   151  // added) to a previously opened zip.Writer.  The archive path of `source` is relative to the
   152  // `rootSource` parameter.
   153  func AddToZip(zipWriter *zip.Writer, source string, rootSource string, logger *logrus.Logger) error {
   154  	return AnnotateAddToZip(zipWriter, source, rootSource, nil, logger)
   155  }