Skip to content

Commit 6a2199d

Browse files
ThesacraftThesacraft
Thesacraft
and
Thesacraft
authored
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 <admin@samuelbellmann.de>
1 parent 0416fde commit 6a2199d

4 files changed

Lines changed: 96 additions & 58 deletions

File tree

‎tagstudio/src/cli/ts_cli.py‎

Lines changed: 34 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
import sys
1313
import time
1414
from PIL import Image, ImageOps, ImageChops, UnidentifiedImageError
15+
from PIL.Image import DecompressionBombError
1516
import pillow_avif
1617
from pathlib import Path
1718
import traceback
@@ -643,8 +644,12 @@ def print_thumbnail(
643644
# raw.thumbnail((512, 512))
644645
raw.thumbnail(self.external_preview_size)
645646
raw.save(external_preview_path)
646-
except:
647-
print(f'{ERROR} Could not load image "{filepath}"')
647+
except (
648+
UnidentifiedImageError,
649+
FileNotFoundError,
650+
DecompressionBombError,
651+
) as e:
652+
print(f'{ERROR} Could not load image "{filepath} due to {e}"')
648653
if self.args.external_preview:
649654
self.set_external_preview_broken()
650655
elif file_type in VIDEO_TYPES:
@@ -1109,24 +1114,34 @@ def create_collage(self) -> str:
11091114
# 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}')
11101115
# sys.stdout.flush()
11111116
if file_type in IMAGE_TYPES:
1112-
with Image.open(
1113-
os.path.normpath(
1114-
f"{self.lib.library_dir}/{entry.path}/{entry.filename}"
1115-
)
1116-
) as pic:
1117-
if keep_aspect:
1118-
pic.thumbnail((thumb_size, thumb_size))
1119-
else:
1120-
pic = pic.resize((thumb_size, thumb_size))
1121-
if data_tint_mode and color:
1122-
pic = pic.convert(mode="RGB")
1123-
pic = ImageChops.hard_light(
1124-
pic,
1125-
Image.new(
1126-
"RGB", (thumb_size, thumb_size), color
1127-
),
1117+
try:
1118+
with Image.open(
1119+
os.path.normpath(
1120+
f"{self.lib.library_dir}/{entry.path}/{entry.filename}"
11281121
)
1129-
collage.paste(pic, (y * thumb_size, x * thumb_size))
1122+
) as pic:
1123+
if keep_aspect:
1124+
pic.thumbnail((thumb_size, thumb_size))
1125+
else:
1126+
pic = pic.resize((thumb_size, thumb_size))
1127+
if data_tint_mode and color:
1128+
pic = pic.convert(mode="RGB")
1129+
pic = ImageChops.hard_light(
1130+
pic,
1131+
Image.new(
1132+
"RGB",
1133+
(thumb_size, thumb_size),
1134+
color,
1135+
),
1136+
)
1137+
collage.paste(
1138+
pic, (y * thumb_size, x * thumb_size)
1139+
)
1140+
except DecompressionBombError as e:
1141+
print(
1142+
f"[ERROR] One of the images was too big ({e})"
1143+
)
1144+
11301145
elif file_type in VIDEO_TYPES:
11311146
video = cv2.VideoCapture(filepath)
11321147
video.set(

‎tagstudio/src/qt/widgets/collage_icon.py‎

Lines changed: 19 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99

1010
import cv2
1111
from PIL import Image, ImageChops, UnidentifiedImageError
12+
from PIL.Image import DecompressionBombError
1213
from PySide6.QtCore import (
1314
QObject,
1415
QThread,
@@ -95,22 +96,25 @@ def render(
9596
# 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}')
9697
# sys.stdout.flush()
9798
if file_type in IMAGE_TYPES:
98-
with Image.open(
99-
os.path.normpath(
100-
f"{self.lib.library_dir}/{entry.path}/{entry.filename}"
101-
)
102-
) as pic:
103-
if keep_aspect:
104-
pic.thumbnail(size)
105-
else:
106-
pic = pic.resize(size)
107-
if data_tint_mode and color:
108-
pic = pic.convert(mode="RGB")
109-
pic = ImageChops.hard_light(
110-
pic, Image.new("RGB", size, color)
99+
try:
100+
with Image.open(
101+
os.path.normpath(
102+
f"{self.lib.library_dir}/{entry.path}/{entry.filename}"
111103
)
112-
# collage.paste(pic, (y*thumb_size, x*thumb_size))
113-
self.rendered.emit(pic)
104+
) as pic:
105+
if keep_aspect:
106+
pic.thumbnail(size)
107+
else:
108+
pic = pic.resize(size)
109+
if data_tint_mode and color:
110+
pic = pic.convert(mode="RGB")
111+
pic = ImageChops.hard_light(
112+
pic, Image.new("RGB", size, color)
113+
)
114+
# collage.paste(pic, (y*thumb_size, x*thumb_size))
115+
self.rendered.emit(pic)
116+
except DecompressionBombError as e:
117+
logging.info(f"[ERROR] One of the images was too big ({e})")
114118
elif file_type in VIDEO_TYPES:
115119
video = cv2.VideoCapture(filepath)
116120
video.set(

‎tagstudio/src/qt/widgets/preview_panel.py‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111

1212
import cv2
1313
from PIL import Image, UnidentifiedImageError
14+
from PIL.Image import DecompressionBombError
1415
from PySide6.QtCore import Signal, Qt, QSize
1516
from PySide6.QtGui import QResizeEvent, QAction
1617
from PySide6.QtWidgets import (
@@ -403,8 +404,15 @@ def update_widgets(self):
403404
)
404405
raise UnidentifiedImageError
405406

406-
except (UnidentifiedImageError, FileNotFoundError, cv2.error):
407-
pass
407+
except (
408+
UnidentifiedImageError,
409+
FileNotFoundError,
410+
cv2.error,
411+
DecompressionBombError,
412+
) as e:
413+
logging.info(
414+
f"[PreviewPanel][ERROR] Couldn't Render thumbnail for {filepath} (because of {e})"
415+
)
408416

409417
try:
410418
self.preview_img.clicked.disconnect()

‎tagstudio/src/qt/widgets/thumb_renderer.py‎

Lines changed: 33 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
ImageOps,
2222
ImageFile,
2323
)
24+
from PIL.Image import DecompressionBombError
2425
from PySide6.QtCore import QObject, Signal, QSize
2526
from PySide6.QtGui import QPixmap
2627
from src.core.ts_core import PLAINTEXT_TYPES, VIDEO_TYPES, IMAGE_TYPES
@@ -138,17 +139,22 @@ def render(
138139
try:
139140
# Images =======================================================
140141
if extension in IMAGE_TYPES:
141-
image = Image.open(filepath)
142-
# image = self.thumb_debug
143-
if image.mode == "RGBA":
144-
# logging.info(image.getchannel(3).tobytes())
145-
new_bg = Image.new("RGB", image.size, color="#1e1e1e")
146-
new_bg.paste(image, mask=image.getchannel(3))
147-
image = new_bg
148-
if image.mode != "RGB":
149-
image = image.convert(mode="RGB")
150-
151-
image = ImageOps.exif_transpose(image)
142+
try:
143+
image = Image.open(filepath)
144+
# image = self.thumb_debug
145+
if image.mode == "RGBA":
146+
# logging.info(image.getchannel(3).tobytes())
147+
new_bg = Image.new("RGB", image.size, color="#1e1e1e")
148+
new_bg.paste(image, mask=image.getchannel(3))
149+
image = new_bg
150+
if image.mode != "RGB":
151+
image = image.convert(mode="RGB")
152+
153+
image = ImageOps.exif_transpose(image)
154+
except DecompressionBombError as e:
155+
logging.info(
156+
f"[ThumbRenderer][ERROR] Couldn't Render thumbnail for {filepath} (because of {e})"
157+
)
152158

153159
# Videos =======================================================
154160
elif extension in VIDEO_TYPES:
@@ -321,17 +327,22 @@ def render_big(
321327
try:
322328
# Images =======================================================
323329
if extension in IMAGE_TYPES:
324-
image = Image.open(filepath)
325-
# image = self.thumb_debug
326-
if image.mode == "RGBA":
327-
# logging.info(image.getchannel(3).tobytes())
328-
new_bg = Image.new("RGB", image.size, color="#1e1e1e")
329-
new_bg.paste(image, mask=image.getchannel(3))
330-
image = new_bg
331-
if image.mode != "RGB":
332-
image = image.convert(mode="RGB")
333-
334-
image = ImageOps.exif_transpose(image)
330+
try:
331+
image = Image.open(filepath)
332+
# image = self.thumb_debug
333+
if image.mode == "RGBA":
334+
# logging.info(image.getchannel(3).tobytes())
335+
new_bg = Image.new("RGB", image.size, color="#1e1e1e")
336+
new_bg.paste(image, mask=image.getchannel(3))
337+
image = new_bg
338+
if image.mode != "RGB":
339+
image = image.convert(mode="RGB")
340+
341+
image = ImageOps.exif_transpose(image)
342+
except DecompressionBombError as e:
343+
logging.info(
344+
f"[ThumbRenderer][ERROR] Couldn't Render thumbnail for {filepath} (because of {e})"
345+
)
335346

336347
# Videos =======================================================
337348
elif extension in VIDEO_TYPES:

0 commit comments

Comments
 (0)