kra/pkg/utils/zip.go

108 lines
2.0 KiB
Go

package utils
import (
"archive/zip"
"io"
"os"
"path/filepath"
"strings"
)
// ZipFiles 压缩文件
func ZipFiles(filename string, files []string, oldForm, newForm string) error {
newZipFile, err := os.Create(filename)
if err != nil {
return err
}
defer newZipFile.Close()
zipWriter := zip.NewWriter(newZipFile)
defer zipWriter.Close()
for _, file := range files {
if err = addFileToZip(zipWriter, file, oldForm, newForm); err != nil {
return err
}
}
return nil
}
func addFileToZip(zipWriter *zip.Writer, filename, oldForm, newForm string) error {
fileToZip, err := os.Open(filename)
if err != nil {
return err
}
defer fileToZip.Close()
info, err := fileToZip.Stat()
if err != nil {
return err
}
header, err := zip.FileInfoHeader(info)
if err != nil {
return err
}
header.Name = strings.Replace(filename, oldForm, newForm, 1)
header.Method = zip.Deflate
writer, err := zipWriter.CreateHeader(header)
if err != nil {
return err
}
_, err = io.Copy(writer, fileToZip)
return err
}
// Unzip 解压文件
func Unzip(src, dest string) ([]string, error) {
var filenames []string
r, err := zip.OpenReader(src)
if err != nil {
return filenames, err
}
defer r.Close()
for _, f := range r.File {
fpath := filepath.Join(dest, f.Name)
// 安全检查:防止路径遍历攻击
if !strings.HasPrefix(fpath, filepath.Clean(dest)+string(os.PathSeparator)) {
continue
}
filenames = append(filenames, fpath)
if f.FileInfo().IsDir() {
os.MkdirAll(fpath, os.ModePerm)
continue
}
if err = os.MkdirAll(filepath.Dir(fpath), os.ModePerm); err != nil {
return filenames, err
}
outFile, err := os.OpenFile(fpath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, f.Mode())
if err != nil {
return filenames, err
}
rc, err := f.Open()
if err != nil {
outFile.Close()
return filenames, err
}
_, err = io.Copy(outFile, rc)
outFile.Close()
rc.Close()
if err != nil {
return filenames, err
}
}
return filenames, nil
}