FranklinWH-Toolbox/service/utils/files/compress.go

package files

import (
	"archive/tar"
	"archive/zip"
	"bufio"
	"bytes"
	"compress/gzip"
	"fmt"
	"io"
	"os"
	"path/filepath"

	"go.uber.org/zap"
)

// ========== ZIP ==========

// 压缩文件或文件夹为 zip
func Zip(srcPath, destZip string) error {
	zipFile, err := os.Create(destZip)
	if err != nil {
		return err
	}
	defer zipFile.Close()

	zw := zip.NewWriter(zipFile)
	defer zw.Close()

	info, err := os.Stat(srcPath)
	if err != nil {
		return err
	}

	if info.IsDir() {
		// 压缩目录
		return filepath.Walk(srcPath, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}

			relPath, err := filepath.Rel(srcPath, path)
			if err != nil {
				return err
			}
			if relPath == "." {
				return nil
			}

			if info.IsDir() {
				_, err := zw.Create(relPath + "/")
				return err
			}

			// 添加文件
			w, err := zw.Create(relPath)
			if err != nil {
				return err
			}
			f, err := os.Open(path)
			if err != nil {
				return err
			}
			defer f.Close()
			_, err = io.Copy(w, f)
			return err
		})
	} else {
		// 压缩单文件
		w, err := zw.Create(info.Name())
		if err != nil {
			return err
		}
		f, err := os.Open(srcPath)
		if err != nil {
			return err
		}
		defer f.Close()

		_, err = io.Copy(w, f)
		return err
	}
}

// 解压 zip
func Unzip(src, dest string) error {
	r, err := zip.OpenReader(src)
	if err != nil {
		return err
	}
	defer r.Close()

	for _, f := range r.File {
		// 拼接目标路径
		fpath := filepath.Join(dest, f.Name)

		// 目录则创建
		if f.FileInfo().IsDir() {
			if err := os.MkdirAll(fpath, os.ModePerm); err != nil {
				return err
			}
			continue
		}

		// 创建父目录
		if err := os.MkdirAll(filepath.Dir(fpath), os.ModePerm); err != nil {
			return err
		}

		// 解压文件
		rc, err := f.Open()
		if err != nil {
			return err
		}
		defer rc.Close()

		out, err := os.Create(fpath)
		if err != nil {
			return err
		}
		defer out.Close()

		if _, err := io.Copy(out, rc); err != nil {
			return err
		}
	}
	return nil
}

// ========== TAR ==========

// 打包文件或目录为 tar
func Tar(srcPath, destTar string) error {
	tarFile, err := os.Create(destTar)
	if err != nil {
		return err
	}
	defer tarFile.Close()

	tw := tar.NewWriter(tarFile)
	defer tw.Close()

	info, err := os.Stat(srcPath)
	if err != nil {
		return err
	}

	if info.IsDir() {
		// 打包目录
		return filepath.Walk(srcPath, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			relPath, err := filepath.Rel(srcPath, path)
			if err != nil {
				return err
			}
			if relPath == "." {
				return nil
			}

			header, err := tar.FileInfoHeader(info, "")
			if err != nil {
				return err
			}
			header.Name = relPath

			if err := tw.WriteHeader(header); err != nil {
				return err
			}

			if !info.Mode().IsRegular() {
				return nil
			}

			f, err := os.Open(path)
			if err != nil {
				return err
			}
			defer f.Close()

			_, err = io.Copy(tw, f)
			return err
		})
	} else {
		// 打包单文件
		header, err := tar.FileInfoHeader(info, "")
		if err != nil {
			return err
		}
		header.Name = info.Name()

		if err := tw.WriteHeader(header); err != nil {
			return err
		}

		if !info.Mode().IsRegular() {
			return nil
		}

		f, err := os.Open(srcPath)
		if err != nil {
			return err
		}
		defer f.Close()

		_, err = io.Copy(tw, f)
		return err
	}
}

// 解包 tar
func Untar(srcTar, destDir string) error {
	f, err := os.Open(srcTar)
	if err != nil {
		return err
	}
	defer f.Close()

	tr := tar.NewReader(f)
	for {
		header, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}

		fpath := filepath.Join(destDir, header.Name)
		switch header.Typeflag {
		case tar.TypeDir:
			if err := os.MkdirAll(fpath, os.ModePerm); err != nil {
				return err
			}
		case tar.TypeReg:
			if err := os.MkdirAll(filepath.Dir(fpath), os.ModePerm); err != nil {
				return err
			}
			out, err := os.Create(fpath)
			if err != nil {
				return err
			}
			defer out.Close()
			if _, err := io.Copy(out, tr); err != nil {
				return err
			}
		}
	}
	return nil
}

// ========== GZIP ==========

// 压缩单文件为 .gz
func Gzip(srcFile, destGz string) error {
	in, err := os.Open(srcFile)
	if err != nil {
		return err
	}
	defer in.Close()

	out, err := os.Create(destGz)
	if err != nil {
		return err
	}
	defer out.Close()

	gw := gzip.NewWriter(out)
	defer gw.Close()

	_, err = io.Copy(gw, in)
	return err
}

// 解压 .gz 为文件
func Gunzip(srcGz, destFile string) error {
	in, err := os.Open(srcGz)
	if err != nil {
		return err
	}
	defer in.Close()

	gr, err := gzip.NewReader(in)
	if err != nil {
		return err
	}
	defer gr.Close()

	out, err := os.Create(destFile)
	if err != nil {
		return err
	}
	defer out.Close()

	_, err = io.Copy(out, gr)
	return err
}

// ========== TGZ ==========

// 压缩文件或文件夹为 tgz
func TarGz(srcPath, destTgz string) error {
	tempFile, err := os.CreateTemp("", "fwhtb-temp-*.tar")
	if err != nil {
		return err
	}
	defer os.Remove(tempFile.Name()) // 程序结束时删除

	// 打包目录为tar文件
	if err = Tar(srcPath, tempFile.Name()); err != nil {
		return err
	}
	// 压缩单tar文件为.tgz或.tar.gz
	if err = Gzip(tempFile.Name(), destTgz); err != nil {
		return err
	}

	return nil
}

// 解压 tgz
func UntarGz(srcTgz, destDir string) error {
	tempFile, err := os.CreateTemp("", "fwhtb-temp-*.tar")
	if err != nil {
		return err
	}
	defer os.Remove(tempFile.Name()) // 程序结束时删除

	// 解压.tgz或.tar.gz为tar文件
	if err = Gunzip(srcTgz, tempFile.Name()); err != nil {
		return err
	}
	// 解包tar文件
	if err = Untar(tempFile.Name(), destDir); err != nil {
		return err
	}

	return nil
}

// ========== 自动解包 ==========

// 自动解包
func AutoExtract(srcFile, destDir string) error {
	ext := filepath.Ext(srcFile)
	name := filepath.Base(srcFile)

	switch {
	// .tar.gz / .tgz
	case ext == ".gz" && (filepath.Ext(name[:len(name)-len(ext)]) == ".tar"):
		return UntarGz(srcFile, destDir)
	case ext == ".tgz":
		return UntarGz(srcFile, destDir)

	// .zip
	case ext == ".zip":
		return Unzip(srcFile, destDir)

	// .tar
	case ext == ".tar":
		return Untar(srcFile, destDir)

	// .gz (单文件压缩)
	case ext == ".gz":
		destFile := filepath.Join(destDir, name[:len(name)-3]) // 去掉.gz
		return Gunzip(srcFile, destFile)

	default:
		// 扩展名不明，尝试魔数检测
		zap.L().Info("未知扩展名，尝试魔数检测")
		return detectByMagic(srcFile, destDir)
	}
}

// 魔数检测
func detectByMagic(srcFile, destDir string) error {
	f, err := os.Open(srcFile)
	if err != nil {
		return err
	}
	defer f.Close()

	reader := bufio.NewReader(f)
	head := make([]byte, 6)
	_, err = reader.Read(head)
	if err != nil {
		return err
	}

	switch {
	case bytes.HasPrefix(head, []byte{0x50, 0x4B, 0x03, 0x04}): // ZIP
		return Unzip(srcFile, destDir)
	case bytes.HasPrefix(head, []byte{0x1F, 0x8B}): // GZIP
		// 可能是 .gz 或 .tar.gz，尝试当 gzip 解开再看
		tempFile := srcFile + ".tmp"
		if err := Gunzip(srcFile, tempFile); err != nil {
			return err
		}
		// 检查是不是 tar
		tf, _ := os.Open(tempFile)
		defer tf.Close()
		buf := make([]byte, 262)
		tf.Read(buf)
		if bytes.Contains(buf[257:262], []byte("ustar")) {
			return UntarGz(srcFile, destDir) // 是 tar.gz
		}
		return Gunzip(srcFile, filepath.Join(destDir, filepath.Base(srcFile[:len(srcFile)-3])))
	}
	return fmt.Errorf("无法识别文件类型: %s", srcFile)
}