archive.go

package utils

import (
	"bytes"
	"compress/bzip2"
	"compress/gzip"
	"io"
	"os"

	"github.com/klauspost/compress/zstd"
	"github.com/ulikunitz/xz"
)

// 压缩格式魔术数
var (
	PrefixZSTD  = []byte{0x28, 0xB5, 0x2F, 0xFD}
	PrefixGZip  = []byte{0x1F, 0x8B}
	PrefixBzip2 = []byte{0x42, 0x5A, 0x68}
	PrefixXZ    = []byte{0xFD, 0x37, 0x7A, 0x58, 0x5A, 0x00}
)

type CompressMethod uint8

const (
	CMUnknown CompressMethod = iota
	CMGZip
	CMBZip2
	CMXZ
	CMZStd
	CMNone // 表示没有压缩
)

// IsTar 判断是否为tar包
func IsTar(r io.Reader) (bool, error) {
	buf := make([]byte, 512)
	rn, err := r.Read(buf)
	if err != nil {
		return false, err
	}
	if rn < 262 {
		return false, nil
	}
	if string(buf[257:262]) == "ustar" {
		return true, nil
	}
	return false, nil
}

// DetectArchive 探测文件是否为压缩包类型
func DetectArchive(path string) (CompressMethod, bool, error) {
	f, err := os.Open(path)
	if err != nil {
		return CMUnknown, false, err
	}
	defer f.Close()
	// 先判断是否为tar包
	istar, err := IsTar(f)
	if err != nil {
		return CMUnknown, false, err
	}
	if istar {
		return CMNone, true, nil
	}
	// 先判断压缩格式
	_, err = f.Seek(0, 0)
	if err != nil {
		return CMUnknown, false, err
	}
	buf := make([]byte, 6)
	rn, err := f.Read(buf)
	if err != nil {
		return CMUnknown, false, err
	}
	if rn != 6 {
		return CMNone, false, nil
	}
	switch {
	case bytes.HasPrefix(buf, PrefixZSTD):
		reader, err := zstd.NewReader(f)
		if err != nil {
			return CMZStd, false, err
		}
		defer reader.Close()
		isTar, err := IsTar(reader)
		return CMZStd, isTar, err
	case bytes.HasPrefix(buf, PrefixGZip):
		reader, err := gzip.NewReader(f)
		if err != nil {
			return CMGZip, false, err
		}
		defer reader.Close()
		isTar, err := IsTar(reader)
		return CMGZip, isTar, err
	case bytes.HasPrefix(buf, PrefixBzip2):
		reader := bzip2.NewReader(f)
		isTar, err := IsTar(reader)
		return CMBZip2, isTar, err
	case bytes.HasPrefix(buf, PrefixXZ):
		reader, err := xz.NewReader(f)
		if err != nil {
			return CMXZ, false, err
		}
		isTar, err := IsTar(reader)
		return CMXZ, isTar, err
	default:
		return CMUnknown, false, nil
	}
}