From 6a2199dd2efb90f4fbd0cb1d9a772eb3f73edbf0 Mon Sep 17 00:00:00 2001 From: Theasacraft <91694323+Thesacraft@users.noreply.github.com> Date: Mon, 13 May 2024 23:18:07 +0200 Subject: [PATCH] Fix pillow decompression bomb error mentioned in #164 (#166) * Fixes DecompressionBombError * Fixes DecompressionBombError in PreviewPanel * Ruff reformat * Handle all DecompressionBombErrors * Handle all DecompressionBombErrors * RUFF * fix typo Co-authored-by: Travis Abendshien <46939827+CyanVoxel@users.noreply.github.com> * fix typo Co-authored-by: Travis Abendshien <46939827+CyanVoxel@users.noreply.github.com> * Ruff reformat --------- Co-authored-by: Thesacraft --- tagstudio/src/cli/ts_cli.py | 53 ++++++++++++++-------- tagstudio/src/qt/widgets/collage_icon.py | 34 ++++++++------ tagstudio/src/qt/widgets/preview_panel.py | 12 ++++- tagstudio/src/qt/widgets/thumb_renderer.py | 51 +++++++++++++-------- 4 files changed, 94 insertions(+), 56 deletions(-) diff --git a/tagstudio/src/cli/ts_cli.py b/tagstudio/src/cli/ts_cli.py index cec3c19d..77626372 100644 --- a/tagstudio/src/cli/ts_cli.py +++ b/tagstudio/src/cli/ts_cli.py @@ -12,6 +12,7 @@ import subprocess import sys import time from PIL import Image, ImageOps, ImageChops, UnidentifiedImageError +from PIL.Image import DecompressionBombError import pillow_avif from pathlib import Path import traceback @@ -643,8 +644,12 @@ class CliDriver: # raw.thumbnail((512, 512)) raw.thumbnail(self.external_preview_size) raw.save(external_preview_path) - except: - print(f'{ERROR} Could not load image "{filepath}"') + except ( + UnidentifiedImageError, + FileNotFoundError, + DecompressionBombError, + ) as e: + print(f'{ERROR} Could not load image "{filepath} due to {e}"') if self.args.external_preview: self.set_external_preview_broken() elif file_type in VIDEO_TYPES: @@ -1109,24 +1114,34 @@ class CliDriver: # sys.stdout.write(f'\r{INFO} Combining [{i+1}/{len(self.lib.entries)}]: {self.get_file_color(file_type)}{entry.path}{os.sep}{entry.filename}{RESET}') # sys.stdout.flush() if file_type in IMAGE_TYPES: - with Image.open( - os.path.normpath( - f"{self.lib.library_dir}/{entry.path}/{entry.filename}" - ) - ) as pic: - if keep_aspect: - pic.thumbnail((thumb_size, thumb_size)) - else: - pic = pic.resize((thumb_size, thumb_size)) - if data_tint_mode and color: - pic = pic.convert(mode="RGB") - pic = ImageChops.hard_light( - pic, - Image.new( - "RGB", (thumb_size, thumb_size), color - ), + try: + with Image.open( + os.path.normpath( + f"{self.lib.library_dir}/{entry.path}/{entry.filename}" ) - collage.paste(pic, (y * thumb_size, x * thumb_size)) + ) as pic: + if keep_aspect: + pic.thumbnail((thumb_size, thumb_size)) + else: + pic = pic.resize((thumb_size, thumb_size)) + if data_tint_mode and color: + pic = pic.convert(mode="RGB") + pic = ImageChops.hard_light( + pic, + Image.new( + "RGB", + (thumb_size, thumb_size), + color, + ), + ) + collage.paste( + pic, (y * thumb_size, x * thumb_size) + ) + except DecompressionBombError as e: + print( + f"[ERROR] One of the images was too big ({e})" + ) + elif file_type in VIDEO_TYPES: video = cv2.VideoCapture(filepath) video.set( diff --git a/tagstudio/src/qt/widgets/collage_icon.py b/tagstudio/src/qt/widgets/collage_icon.py index b8380a3c..d83fc958 100644 --- a/tagstudio/src/qt/widgets/collage_icon.py +++ b/tagstudio/src/qt/widgets/collage_icon.py @@ -9,6 +9,7 @@ from pathlib import Path import cv2 from PIL import Image, ImageChops, UnidentifiedImageError +from PIL.Image import DecompressionBombError from PySide6.QtCore import ( QObject, QThread, @@ -95,22 +96,25 @@ class CollageIconRenderer(QObject): # sys.stdout.write(f'\r{INFO} Combining [{i+1}/{len(self.lib.entries)}]: {self.get_file_color(file_type)}{entry.path}{os.sep}{entry.filename}{RESET}') # sys.stdout.flush() if file_type in IMAGE_TYPES: - with Image.open( - os.path.normpath( - f"{self.lib.library_dir}/{entry.path}/{entry.filename}" - ) - ) as pic: - if keep_aspect: - pic.thumbnail(size) - else: - pic = pic.resize(size) - if data_tint_mode and color: - pic = pic.convert(mode="RGB") - pic = ImageChops.hard_light( - pic, Image.new("RGB", size, color) + try: + with Image.open( + os.path.normpath( + f"{self.lib.library_dir}/{entry.path}/{entry.filename}" ) - # collage.paste(pic, (y*thumb_size, x*thumb_size)) - self.rendered.emit(pic) + ) as pic: + if keep_aspect: + pic.thumbnail(size) + else: + pic = pic.resize(size) + if data_tint_mode and color: + pic = pic.convert(mode="RGB") + pic = ImageChops.hard_light( + pic, Image.new("RGB", size, color) + ) + # collage.paste(pic, (y*thumb_size, x*thumb_size)) + self.rendered.emit(pic) + except DecompressionBombError as e: + logging.info(f"[ERROR] One of the images was too big ({e})") elif file_type in VIDEO_TYPES: video = cv2.VideoCapture(filepath) video.set( diff --git a/tagstudio/src/qt/widgets/preview_panel.py b/tagstudio/src/qt/widgets/preview_panel.py index a4b2eb7d..81d8c97a 100644 --- a/tagstudio/src/qt/widgets/preview_panel.py +++ b/tagstudio/src/qt/widgets/preview_panel.py @@ -11,6 +11,7 @@ from datetime import datetime as dt import cv2 from PIL import Image, UnidentifiedImageError +from PIL.Image import DecompressionBombError from PySide6.QtCore import Signal, Qt, QSize from PySide6.QtGui import QResizeEvent, QAction from PySide6.QtWidgets import ( @@ -403,8 +404,15 @@ class PreviewPanel(QWidget): ) raise UnidentifiedImageError - except (UnidentifiedImageError, FileNotFoundError, cv2.error): - pass + except ( + UnidentifiedImageError, + FileNotFoundError, + cv2.error, + DecompressionBombError, + ) as e: + logging.info( + f"[PreviewPanel][ERROR] Couldn't Render thumbnail for {filepath} (because of {e})" + ) try: self.preview_img.clicked.disconnect() diff --git a/tagstudio/src/qt/widgets/thumb_renderer.py b/tagstudio/src/qt/widgets/thumb_renderer.py index c3e6bbff..5e87d7a0 100644 --- a/tagstudio/src/qt/widgets/thumb_renderer.py +++ b/tagstudio/src/qt/widgets/thumb_renderer.py @@ -21,6 +21,7 @@ from PIL import ( ImageOps, ImageFile, ) +from PIL.Image import DecompressionBombError from PySide6.QtCore import QObject, Signal, QSize from PySide6.QtGui import QPixmap from src.core.ts_core import PLAINTEXT_TYPES, VIDEO_TYPES, IMAGE_TYPES @@ -138,17 +139,22 @@ class ThumbRenderer(QObject): try: # Images ======================================================= if extension in IMAGE_TYPES: - image = Image.open(filepath) - # image = self.thumb_debug - if image.mode == "RGBA": - # logging.info(image.getchannel(3).tobytes()) - new_bg = Image.new("RGB", image.size, color="#1e1e1e") - new_bg.paste(image, mask=image.getchannel(3)) - image = new_bg - if image.mode != "RGB": - image = image.convert(mode="RGB") + try: + image = Image.open(filepath) + # image = self.thumb_debug + if image.mode == "RGBA": + # logging.info(image.getchannel(3).tobytes()) + new_bg = Image.new("RGB", image.size, color="#1e1e1e") + new_bg.paste(image, mask=image.getchannel(3)) + image = new_bg + if image.mode != "RGB": + image = image.convert(mode="RGB") - image = ImageOps.exif_transpose(image) + image = ImageOps.exif_transpose(image) + except DecompressionBombError as e: + logging.info( + f"[ThumbRenderer][ERROR] Couldn't Render thumbnail for {filepath} (because of {e})" + ) # Videos ======================================================= elif extension in VIDEO_TYPES: @@ -321,17 +327,22 @@ class ThumbRenderer(QObject): try: # Images ======================================================= if extension in IMAGE_TYPES: - image = Image.open(filepath) - # image = self.thumb_debug - if image.mode == "RGBA": - # logging.info(image.getchannel(3).tobytes()) - new_bg = Image.new("RGB", image.size, color="#1e1e1e") - new_bg.paste(image, mask=image.getchannel(3)) - image = new_bg - if image.mode != "RGB": - image = image.convert(mode="RGB") + try: + image = Image.open(filepath) + # image = self.thumb_debug + if image.mode == "RGBA": + # logging.info(image.getchannel(3).tobytes()) + new_bg = Image.new("RGB", image.size, color="#1e1e1e") + new_bg.paste(image, mask=image.getchannel(3)) + image = new_bg + if image.mode != "RGB": + image = image.convert(mode="RGB") - image = ImageOps.exif_transpose(image) + image = ImageOps.exif_transpose(image) + except DecompressionBombError as e: + logging.info( + f"[ThumbRenderer][ERROR] Couldn't Render thumbnail for {filepath} (because of {e})" + ) # Videos ======================================================= elif extension in VIDEO_TYPES: