star/mime/mime.go

package mime

import (
	"b612.me/starlog"
	"b612.me/staros"
	"encoding/csv"
	"github.com/gabriel-vasile/mimetype"
	"github.com/spf13/cobra"
	"io"
	"os"
	"path/filepath"
	"strings"
)

var recursive, verbose, repair, appendOldName, onlyShowNotMatch, secondExt bool
var outputCsvPath string
var fromSize int

func init() {
	Cmd.Flags().BoolVarP(&recursive, "recursive", "r", false, "递归检测文件夹中的所有文件")
	Cmd.Flags().BoolVarP(&verbose, "verbose", "v", false, "详细模式，输出更多日志")
	Cmd.Flags().BoolVarP(&repair, "repair", "R", false, "修复文件类型，尝试将文件类型修复为正确的类型")
	Cmd.Flags().BoolVarP(&appendOldName, "append-mode", "a", false, "在修复文件名时，追加旧文件名后缀")
	Cmd.Flags().StringVarP(&outputCsvPath, "output-csv", "o", "", "输出CSV文件路径，若不指定则不输出CSV")
	Cmd.Flags().IntVarP(&fromSize, "from-size", "s", 0, "需要检测的地址偏移量，单位为字节，默认从0开始检测，负数表示从文件末尾开始计算偏移量")
	Cmd.Flags().BoolVarP(&onlyShowNotMatch, "only-show-not-match", "m", false, "只显示文件类型不匹配的文件")
	Cmd.Flags().BoolVarP(&secondExt, "second-ext", "S", false, "使用第二个扩展名进行匹配，适用于某些文件类型（如 .tar.gz .7z.001）")
}

var Cmd = &cobra.Command{
	Use:   "mime",
	Short: "文件类型检测",
	Long:  `mime - 文件类型检测`,
	Run: func(cmd *cobra.Command, args []string) {
		var c *csv.Writer
		if len(args) == 0 {
			starlog.Errorf("请指定文件路径\n")
			os.Exit(1)
		}
		if outputCsvPath != "" {
			file, err := os.Create(outputCsvPath)
			if err != nil {
				starlog.Errorf("创建CSV文件 %s 失败: %v\n", outputCsvPath, err)
				os.Exit(1)
			}
			defer file.Close()
			c = csv.NewWriter(file)
			c.Write([]string{"File Path", "Extension", "Detect File Type", "Detect File Extension", "Match"})
			defer c.Flush()
		}
		for _, path := range args {
			if staros.IsFolder(path) {
				for _, file := range getFiles(path, recursive) {
					Detect(c, file, fromSize, onlyShowNotMatch)
				}
			} else {
				Detect(c, path, fromSize, onlyShowNotMatch)
			}
		}
	},
}

func getFiles(folder string, recurring bool) []string {
	files := make([]string, 0)
	err := filepath.Walk(folder, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			starlog.Errorf("walk folder %s error: %v", folder, err)
			return err
		}
		if info.IsDir() {
			if recurring {
				return nil // 继续递归
			}
			return filepath.SkipDir // 不递归子目录
		}
		files = append(files, path)
		return nil
	})
	if err != nil {
		starlog.Errorf("walk folder %s error: %v", folder, err)
	}
	return files
}

func Detect(c *csv.Writer, path string, fromSize int, onlyShowNotMatch bool) {
	f, err := os.Open(path)
	if err != nil {
		starlog.Errorf("open file %s error: %v\n", path, err)
		os.Exit(1)
	}
	if fromSize > 0 {
		f.Seek(int64(fromSize), io.SeekStart)
	} else if fromSize < 0 {
		f.Seek(int64(fromSize), io.SeekEnd)
	}
	tp, err := mimetype.DetectReader(f)
	f.Close()
	if err != nil {
		starlog.Errorf("detect file type error: %v\n", err)
		os.Exit(1)
	}
	ext := strings.ToLower(filepath.Ext(path))
	ext2 := strings.ToLower(filepath.Ext(strings.TrimRight(path, ext)))
	if verbose {
		starlog.Infof("file path: %s\n", path)
		starlog.Infof("file type: %s\n", tp.String())
		starlog.Infof("file ext: %s\n", tp.Extension())
		if ext2 == tp.Extension() && secondExt {
			if !onlyShowNotMatch {
				starlog.Infof("file type %s match extension %s\n", tp.String(), ext2)
			}
		} else if ext != tp.Extension() && tp.Extension() != "" {
			starlog.Warningf("file type %s not match extension %s\n", tp.String(), ext)
		} else if !onlyShowNotMatch {
			starlog.Infof("file type %s match extension %s\n", tp.String(), ext)
		}
	} else {
		if ext2 == tp.Extension() && secondExt {
			if !onlyShowNotMatch {
				starlog.Infof("%s: %s (%s) matched %s\n", path, tp.String(), tp.Extension(), ext2)
			}
		} else if ext != tp.Extension() && tp.Extension() != "" {
			starlog.Warningf("%s: %s (%s) not match %s\n", path, tp.String(), tp.Extension(), ext)
		} else if !onlyShowNotMatch {
			starlog.Infof("%s: %s (%s) matched %s\n", path, tp.String(), tp.Extension(), ext)
		}
	}
	if c != nil {
		match := "false"
		if ext == tp.Extension() {
			match = "true"
		}
		err = c.Write([]string{path, ext, tp.String(), tp.Extension(), match})
		if err != nil {
			starlog.Errorf("write csv file %s error: %v\n", outputCsvPath, err)
		}
	}
	if repair && ext != tp.Extension() && tp.Extension() != "" {
		newPath := path[:len(path)-len(ext)] + tp.Extension()
		if appendOldName {
			newPath = path + tp.Extension()
		}
		err = os.Rename(path, newPath)
		if err != nil {
			starlog.Errorf("rename file %s to %s error: %v", path, newPath, err)
			os.Exit(1)
		}
		if verbose {
			starlog.Infof("file %s renamed to %s\n", path, newPath)
		} else {
			starlog.Infof("renamed: %s -> %s\n", path, newPath)
		}
	}
}