goext/zipext/zip.go

183 lines
2.6 KiB
Go
Raw Permalink Normal View History

2022-10-27 16:00:57 +02:00
package zipext
import (
"archive/tar"
"archive/zip"
"bufio"
"bytes"
"compress/flate"
"compress/gzip"
"errors"
"time"
)
var errAlreadyClosed = errors.New("already closed")
var errZipNotEnabled = errors.New("zip not enabled")
var errTgzNotEnabled = errors.New("tgz not enabled")
type MemoryZip struct {
zipEnabled bool
zipbuffer *bytes.Buffer
zipwriter *zip.Writer
tarEnabled bool
tarbuffer *bytes.Buffer
tarwriter *tar.Writer
open bool
}
func NewMemoryZip(enableGZip, enableTarGZ bool) *MemoryZip {
var bz *bytes.Buffer = nil
var z *zip.Writer = nil
var bt *bytes.Buffer = nil
var t *tar.Writer = nil
if enableGZip {
bz = new(bytes.Buffer)
z = zip.NewWriter(bz)
}
if enableTarGZ {
bt = new(bytes.Buffer)
t = tar.NewWriter(bt)
}
return &MemoryZip{
zipEnabled: enableGZip,
zipbuffer: bz,
zipwriter: z,
tarEnabled: enableTarGZ,
tarbuffer: bz,
tarwriter: t,
open: true,
}
}
func (z *MemoryZip) AddFile(path string, data []byte) error {
var err error
if !z.open {
return errAlreadyClosed
}
if z.zipEnabled {
zipheader, err := z.zipwriter.CreateHeader(&zip.FileHeader{
Name: path,
Method: zip.Deflate,
Modified: time.Now(),
})
if err != nil {
return err
}
_, err = zipheader.Write(data)
if err != nil {
return err
}
}
if z.tarEnabled {
tarheader := &tar.Header{
Name: path,
ModTime: time.Now(),
Typeflag: tar.TypeReg,
Size: int64(len(data)),
}
err = z.tarwriter.WriteHeader(tarheader)
if err != nil {
return err
}
_, err = z.tarwriter.Write(data)
if err != nil {
return err
}
}
return nil
}
func (z *MemoryZip) GetZip() ([]byte, error) {
if !z.zipEnabled {
return nil, errZipNotEnabled
}
if z.open {
err := z.Close()
if err != nil {
return nil, err
}
}
return z.zipbuffer.Bytes(), nil
}
func (z *MemoryZip) GetTarGz() ([]byte, error) {
if !z.tarEnabled {
return nil, errTgzNotEnabled
}
if z.open {
err := z.Close()
if err != nil {
return nil, err
}
}
b := new(bytes.Buffer)
gf, err := gzip.NewWriterLevel(b, flate.BestCompression)
if err != nil {
return nil, err
}
fw := bufio.NewWriter(gf)
_, err = fw.Write(z.tarbuffer.Bytes())
if err != nil {
return nil, err
}
err = fw.Flush()
if err != nil {
return nil, err
}
err = gf.Close()
if err != nil {
return nil, err
}
return b.Bytes(), nil
}
func (z *MemoryZip) Close() error {
if !z.open {
return nil
}
z.open = false
if z.zipEnabled {
err := z.zipwriter.Close()
if err != nil {
return err
}
}
if z.tarEnabled {
err := z.tarwriter.Close()
if err != nil {
return err
}
}
return nil
}