145 lines
3.8 KiB
Go
145 lines
3.8 KiB
Go
package extractor
|
|
|
|
import (
|
|
"archive/tar"
|
|
"compress/gzip"
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
|
|
"yizhisec.com/hsv2/forge/pkg/logger"
|
|
)
|
|
|
|
type Options struct {
|
|
OnProgress func(filename string, index int, total int)
|
|
IsGzipped *bool
|
|
}
|
|
|
|
type Option func(*Options)
|
|
|
|
func WithProgress(callback func(filename string, index int, total int)) Option {
|
|
return func(o *Options) {
|
|
o.OnProgress = callback
|
|
}
|
|
}
|
|
|
|
func WithGzipCompression(isGzipped bool) Option {
|
|
return func(o *Options) {
|
|
o.IsGzipped = &isGzipped
|
|
}
|
|
}
|
|
|
|
func Extract(ctx context.Context, src, destDir string, opts ...Option) error {
|
|
options := &Options{}
|
|
for _, opt := range opts {
|
|
opt(options)
|
|
}
|
|
|
|
logger.Debug("开始解压: %s -> %s", src, destDir)
|
|
|
|
f, err := os.Open(src)
|
|
if err != nil {
|
|
logger.Debug("打开源文件失败 %s: %v", src, err)
|
|
return fmt.Errorf("failed to open source file: %w", err)
|
|
}
|
|
defer f.Close()
|
|
|
|
var reader io.Reader = f
|
|
var isGzipped bool
|
|
|
|
if options.IsGzipped != nil {
|
|
isGzipped = *options.IsGzipped
|
|
} else {
|
|
isGzipped = strings.HasSuffix(strings.ToLower(src), ".tar.gz") || strings.HasSuffix(strings.ToLower(src), ".tgz")
|
|
}
|
|
|
|
if isGzipped {
|
|
gzReader, err := gzip.NewReader(f)
|
|
if err != nil {
|
|
logger.Debug("创建 gzip reader 失败: %v", err)
|
|
return fmt.Errorf("failed to create gzip reader: %w", err)
|
|
}
|
|
defer gzReader.Close()
|
|
reader = gzReader
|
|
}
|
|
|
|
if err := extractTar(ctx, reader, destDir, options); err != nil {
|
|
return err
|
|
}
|
|
|
|
logger.Debug("解压完成: %s", src)
|
|
return nil
|
|
}
|
|
|
|
func extractTar(ctx context.Context, r io.Reader, destDir string, options *Options) error {
|
|
tarReader := tar.NewReader(r)
|
|
fileCount := 0
|
|
totalFiles := 0
|
|
|
|
for {
|
|
if ctx.Err() != nil {
|
|
return ctx.Err()
|
|
}
|
|
|
|
header, err := tarReader.Next()
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
if err != nil {
|
|
logger.Debug("读取 tar 条目失败: %v", err)
|
|
return fmt.Errorf("failed to read tar header: %w", err)
|
|
}
|
|
|
|
target := filepath.Join(destDir, header.Name)
|
|
logger.Debug("解压文件: %s", header.Name)
|
|
|
|
if options.OnProgress != nil {
|
|
options.OnProgress(header.Name, fileCount, totalFiles)
|
|
}
|
|
|
|
switch header.Typeflag {
|
|
case tar.TypeDir:
|
|
if err := os.MkdirAll(target, os.FileMode(header.Mode)); err != nil {
|
|
logger.Debug("创建目录失败 %s: %v", target, err)
|
|
return fmt.Errorf("failed to create directory %s: %w", target, err)
|
|
}
|
|
case tar.TypeReg:
|
|
if err := os.MkdirAll(filepath.Dir(target), 0755); err != nil {
|
|
logger.Debug("创建父目录失败 %s: %v", filepath.Dir(target), err)
|
|
return fmt.Errorf("failed to create parent directory: %w", err)
|
|
}
|
|
outFile, err := os.OpenFile(target, os.O_CREATE|os.O_RDWR|os.O_TRUNC, os.FileMode(header.Mode))
|
|
if err != nil {
|
|
logger.Debug("创建文件失败 %s: %v", target, err)
|
|
return fmt.Errorf("failed to create file %s: %w", target, err)
|
|
}
|
|
if _, err := io.Copy(outFile, tarReader); err != nil {
|
|
outFile.Close()
|
|
logger.Debug("写入文件失败 %s: %v", target, err)
|
|
return fmt.Errorf("failed to write file %s: %w", target, err)
|
|
}
|
|
outFile.Close()
|
|
fileCount++
|
|
case tar.TypeSymlink:
|
|
if err := os.MkdirAll(filepath.Dir(target), 0755); err != nil {
|
|
logger.Debug("创建符号链接父目录失败 %s: %v", filepath.Dir(target), err)
|
|
return fmt.Errorf("failed to create parent directory for symlink: %w", err)
|
|
}
|
|
os.Remove(target)
|
|
if err := os.Symlink(header.Linkname, target); err != nil {
|
|
logger.Debug("创建符号链接失败 %s -> %s: %v", target, header.Linkname, err)
|
|
return fmt.Errorf("failed to create symlink %s -> %s: %w", target, header.Linkname, err)
|
|
}
|
|
fileCount++
|
|
default:
|
|
logger.Debug("跳过不支持的文件类型: %s (type: %v)", header.Name, header.Typeflag)
|
|
}
|
|
}
|
|
|
|
logger.Debug("解压完成，共 %d 个文件", fileCount)
|
|
return nil
|
|
}
|