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