From 0849c99d3c31697e85745568f61e3c6352c683b2 Mon Sep 17 00:00:00 2001
From: blank X <blankie@nixnetmail.com>
Date: Thu, 18 Nov 2021 17:44:24 +0700
Subject: [PATCH] more fixes idk this history gonna be public n fun

---
 main.go    |  9 +++++----
 structs.go |  8 ++++++--
 utils.go   | 35 ++++++++++++++++++++++++++++++++---
 3 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/main.go b/main.go
index 954f381..9075a5a 100644
--- a/main.go
+++ b/main.go
@@ -21,6 +21,7 @@ func main() {
 		fmt.Fprintf(os.Stderr, "Usage: %s <submission id/url>\n", os.Args[0])
 		os.Exit(1)
 	}
+	LoadMimetypes()
 	submissionUrl, err := url.Parse(os.Args[1])
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "Error when parsing submission url: %s\n", err)
@@ -196,10 +197,10 @@ func main() {
 	} else if submission.IsRedditMediaDomain && submission.Preview != nil {
 		preview := submission.Preview.Images[0]
 		if pathExt == ".gif" {
-			if preview.Variants.Mp4 != nil {
-				urls[0] = preview.Variants.Mp4.Url
-			} else if preview.Variants.Gif != nil {
-				urls[0] = preview.Variants.Gif.Url
+			if preview.Variants.Mp4.Source != nil {
+				urls[0] = preview.Variants.Mp4.Source.Url
+			} else if preview.Variants.Gif.Source != nil {
+				urls[0] = preview.Variants.Gif.Source.Url
 			}
 		} else if submission.IsVideo {
 			urls[0] = preview.Source.Url
diff --git a/structs.go b/structs.go
index e289e2c..a82d48d 100644
--- a/structs.go
+++ b/structs.go
@@ -64,8 +64,12 @@ type Submission struct {
 	Preview             *struct {
 		Images []struct {
 			Variants struct {
-				Mp4 *PreviewSource `json:"mp4"`
-				Gif *PreviewSource `json:"gif"`
+				Mp4 struct {
+					Source *PreviewSource `json:"source"`
+				} `json:"mp4"`
+				Gif struct {
+					Source *PreviewSource `json:"source"`
+				} `json:"gif"`
 			} `json:"variants"`
 			Source *PreviewSource `json:"source"`
 		} `json:"images"`
diff --git a/utils.go b/utils.go
index 496872b..b94fc80 100644
--- a/utils.go
+++ b/utils.go
@@ -17,9 +17,10 @@ import (
 )
 
 var (
-	configDir string
-	config    *Config
-	data      *Data
+	configDir    string
+	config       *Config
+	data         *Data
+	mimeOverride map[string]string
 )
 
 func LoadConfigAndData() error {
@@ -74,6 +75,30 @@ func LoadConfigAndData() error {
 	return nil
 }
 
+// a jpeg file can have a jpe extension and i personally don't like it
+// https://github.com/LonamiWebs/Telethon/blob/2e1be01ad4f6462de2e9e1f96a33537e51f44980/telethon/utils.py#L33
+func LoadMimetypes() {
+	mimeOverride = make(map[string]string)
+	mimeOverride["image/png"] = ".png"
+	mimeOverride["image/jpeg"] = ".jpeg"
+	mimeOverride["image/webp"] = ".webp"
+	mimeOverride["image/gif"] = ".gif"
+	mimeOverride["image/bmp"] = ".bmp"
+	mimeOverride["image/x-tga"] = ".tga"
+	mimeOverride["image/tiff"] = ".tiff"
+	mimeOverride["image/vnd.adobe.photoshop"] = ".psd"
+
+	mimeOverride["video/mp4"] = ".mp4"
+	mimeOverride["video/quicktime"] = ".mov"
+	mimeOverride["video/avi"] = ".avi"
+
+	mimeOverride["audio/mpeg"] = ".mp3"
+	mimeOverride["audio/m4a"] = ".m4a"
+	mimeOverride["audio/aac"] = ".aac"
+	mimeOverride["audio/ogg"] = ".ogg"
+	mimeOverride["audio/flac"] = ".flac"
+}
+
 func WriteData() error {
 	contents, err := json.Marshal(data)
 	if err != nil {
@@ -256,6 +281,10 @@ func GetExtension(mimeType string) (string, error) {
 	if err != nil {
 		return "", fmt.Errorf("Failed to parse media type of %s: %s", mimeType, err)
 	}
+	ext, exists := mimeOverride[mediaType]
+	if exists {
+		return ext, nil
+	}
 	exts, err := mime.ExtensionsByType(mediaType)
 	if err != nil {
 		return "", fmt.Errorf("Failed to find file extensions of %s: %s", mediaType, err)