Support fixing more extensions than just jpeg -> png

main.go

@@ -9,6 +9,7 @@ import (
 	"os"
 	"path"
 	"path/filepath"
+	"slices"
 	"strings"
 
 	"github.com/gabriel-vasile/mimetype"
@@ -19,6 +20,23 @@ type Config struct {
 	Paths []string
 }
 
+// Define a map of MIME types to their correct extensions.
+var conversionMap = map[string][]string{
+	"image/png":              {".png"},
+	"image/vnd.mozilla.apng": {".png"},
+	"image/jpeg":             {".jpg", ".jpeg"},
+	"image/webp":             {".webp"},
+	"video/mp4":              {".mp4"},
+	"video/quicktime":        {".mp4"},
+	"video/3gpp":             {".3gp", ".3gpp", ".mp4"},
+	"video/x-m4v":            {".m4v", ".mp4"},
+	"video/webm":             {".webm"},
+	"video/x-matroska":       {".mkv"},
+	"image/gif":              {".gif"},
+}
+
+var suffixesToCheck []string
+
 func loadConfiguration() (*Config, error) {
 	dir, err := os.UserConfigDir()
 	if err != nil {
@@ -67,7 +85,19 @@ func loadConfiguration() (*Config, error) {
 	return &config, nil
 }
 
+func loadSuffixes() {
+	for _, extensions := range conversionMap {
+		// Only add unique extensions to the suffixesToCheck list.
+		for _, ext := range extensions {
+			if !slices.Contains(suffixesToCheck, ext) {
+				suffixesToCheck = append(suffixesToCheck, ext)
+			}
+		}
+	}
+}
+
 func main() {
+	loadSuffixes()
 	config, err := loadConfiguration()
 	if err != nil {
 		log.Fatal("Failed to load configuration:", err)
@@ -82,13 +112,17 @@ func main() {
 	}
 }
 
+// This function returns true if the entry is a regular file and has an extension that is in the suffixesToCheck list.
 func shouldEntryBeChecked(entry fs.DirEntry) bool {
-	//fmt.Println(filepath.Ext(entry.Name()))
-	// If it's a file
-	return (entry.Type().IsRegular() &&
-		// If it ends with .jpeg or .jpg
-		(strings.HasSuffix(entry.Name(), ".jpeg") ||
-			strings.HasSuffix(entry.Name(), ".jpg")))
+	if !entry.Type().IsRegular() {
+		return false
+	}
+	for _, suffix := range suffixesToCheck {
+		if strings.HasSuffix(entry.Name(), suffix) {
+			return true
+		}
+	}
+	return false
 }
 
 func getFileMimeType(pathToFile string) *mimetype.MIME {
@@ -105,6 +139,10 @@ func changeFileExtension(originalPath string, newExtension string) {
 	os.Rename(originalPath, newPath)
 }
 
+func isFileExtensionExpected(currentExt string, expectedExts []string) bool {
+	return slices.Contains(expectedExts, currentExt)
+}
+
 func fixExtensionsInDirectory(directory string) {
 	entries, err := os.ReadDir(directory)
 	if err != nil {
@@ -115,9 +153,18 @@ func fixExtensionsInDirectory(directory string) {
 	for _, v := range entries {
 		if shouldEntryBeChecked(v) {
 			path := filepath.Join(directory, v.Name())
-			if getFileMimeType(path).String() == "image/png" {
-				fmt.Println("File at", path, "was detected as png with wrong extension")
-				changeFileExtension(path, ".png")
+			fileExt := filepath.Ext(path)
+			mimeType := getFileMimeType(path).String()
+
+			expectedExtensions, ok := conversionMap[mimeType]
+			if !ok {
+				fmt.Println("File at", path, "has unrecognized MIME type", mimeType)
+				continue
+			}
+			// Check if the file's extension is in the list of expected extensions for its MIME type.
+			if !isFileExtensionExpected(fileExt, expectedExtensions) {
+				fmt.Println("File at", path, "has MIME", mimeType, "and expects", expectedExtensions, "but got", fileExt)
+				changeFileExtension(path, expectedExtensions[0])
 			}
 		} else if v.IsDir() {
 			fixExtensionsInDirectory(path.Join(directory, v.Name()))