From f2a4ea1794718e4dc0148bc172cb877f1080903b Mon Sep 17 00:00:00 2001
From: Simon Sawicki <contact@grub4k.xyz>
Date: Thu, 27 Jun 2024 16:12:19 +0200
Subject: [PATCH] [pp/embedthumbnail] Fix postprocessor (#10248)

* [compat] Improve `imghdr.what` detection
* [pp/embedthumbnail] Improve imghdr fail message
* [pp/embedthumbnail] Fix AtomicParsley error handling

Authored by: Grub4K
---
 yt_dlp/compat/imghdr.py                | 26 ++++++++++++++++----------
 yt_dlp/postprocessor/embedthumbnail.py | 18 +++++++++++++-----
 2 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/yt_dlp/compat/imghdr.py b/yt_dlp/compat/imghdr.py
index 5d64ab07b..4ae173fde 100644
--- a/yt_dlp/compat/imghdr.py
+++ b/yt_dlp/compat/imghdr.py
@@ -1,16 +1,22 @@
-tests = {
-    'webp': lambda h: h[0:4] == b'RIFF' and h[8:] == b'WEBP',
-    'png': lambda h: h[:8] == b'\211PNG\r\n\032\n',
-    'jpeg': lambda h: h[6:10] in (b'JFIF', b'Exif'),
-    'gif': lambda h: h[:6] in (b'GIF87a', b'GIF89a'),
-}
-
-
 def what(file=None, h=None):
     """Detect format of image (Currently supports jpeg, png, webp, gif only)
-    Ref: https://github.com/python/cpython/blob/3.10/Lib/imghdr.py
+    Ref: https://github.com/python/cpython/blob/3.11/Lib/imghdr.py
+    Ref: https://www.w3.org/Graphics/JPEG/itu-t81.pdf
     """
     if h is None:
         with open(file, 'rb') as f:
             h = f.read(12)
-    return next((type_ for type_, test in tests.items() if test(h)), None)
+
+    if h.startswith(b'RIFF') and h.startswith(b'WEBP', 8):
+        return 'webp'
+
+    if h.startswith(b'\x89PNG'):
+        return 'png'
+
+    if h.startswith(b'\xFF\xD8\xFF'):
+        return 'jpeg'
+
+    if h.startswith(b'GIF'):
+        return 'gif'
+
+    return None
diff --git a/yt_dlp/postprocessor/embedthumbnail.py b/yt_dlp/postprocessor/embedthumbnail.py
index 673a92468..f2228ac61 100644
--- a/yt_dlp/postprocessor/embedthumbnail.py
+++ b/yt_dlp/postprocessor/embedthumbnail.py
@@ -119,14 +119,21 @@ def run(self, info):
             if not mutagen or prefer_atomicparsley:
                 success = False
             else:
+                self._report_run('mutagen', filename)
+                f = {'jpeg': MP4Cover.FORMAT_JPEG, 'png': MP4Cover.FORMAT_PNG}
                 try:
-                    self._report_run('mutagen', filename)
+                    with open(thumbnail_filename, 'rb') as thumbfile:
+                        thumb_data = thumbfile.read()
+
+                    type_ = imghdr.what(h=thumb_data)
+                    if not type_:
+                        raise ValueError('could not determine image type')
+                    elif type_ not in f:
+                        raise ValueError(f'incompatible image type: {type_}')
+
                     meta = MP4(filename)
                     # NOTE: the 'covr' atom is a non-standard MPEG-4 atom,
                     # Apple iTunes 'M4A' files include the 'moov.udta.meta.ilst' atom.
-                    f = {'jpeg': MP4Cover.FORMAT_JPEG, 'png': MP4Cover.FORMAT_PNG}[imghdr.what(thumbnail_filename)]
-                    with open(thumbnail_filename, 'rb') as thumbfile:
-                        thumb_data = thumbfile.read()
                     meta.tags['covr'] = [MP4Cover(data=thumb_data, imageformat=f)]
                     meta.save()
                     temp_filename = filename
@@ -160,9 +167,10 @@ def run(self, info):
                     stdout, stderr, returncode = Popen.run(cmd, text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                     if returncode:
                         self.report_warning(f'Unable to embed thumbnails using AtomicParsley; {stderr.strip()}')
+                        success = False
                     # for formats that don't support thumbnails (like 3gp) AtomicParsley
                     # won't create to the temporary file
-                    if 'No changes' in stdout:
+                    elif 'No changes' in stdout:
                         self.report_warning('The file format doesn\'t support embedding a thumbnail')
                         success = False