import contextlib import logging import os from collections.abc import Iterator from typing import Any import boto3 import botocore import pyvips from botocore.client import Config from django.conf import settings from django.db import migrations from django.db.backends.base.schema import BaseDatabaseSchemaEditor from django.db.migrations.state import StateApps DEFAULT_EMOJI_SIZE = 64 IMAGE_BOMB_TOTAL_PIXELS = 90000000 class SkipImageError(Exception): pass # From zerver.lib.thumbnail, with minor exception changes @contextlib.contextmanager def libvips_check_image(image_data: bytes) -> Iterator[pyvips.Image]: try: source_image = pyvips.Image.new_from_buffer(image_data, "") except pyvips.Error as e: raise SkipImageError(f"Cannot process image: {e}") if ( source_image.width * source_image.height * source_image.get_n_pages() > IMAGE_BOMB_TOTAL_PIXELS ): raise SkipImageError("Image size exceeds limit.") try: yield source_image except pyvips.Error as e: raise SkipImageError(f"Bad image data? {e}") def extract_still_from_animated( emoji: Any, image_data: bytes, maybe_skip: bool = False, size: int = DEFAULT_EMOJI_SIZE ) -> bytes | None: # Decide whether we need to (re)generate the still frame: # # 1. If the still_frame is not square, we must always pad and regenerate, # even if a still already exists. # # 2. If maybe_skip is True, we have the original file available and the # emoji is already marked animated. In this case, a correct still was # previously generated and the image is already square, so we can skip # regeneration(We always try to load the original file, but in the # unlikely case that it isn’t available, we load the resized file instead). # # 3. If emoji.is_animated is False, we always generate a still. # # 4. If we don't have the original file (maybe_skip is False), we always # regenerate since we can't trust any previously generated still. with libvips_check_image(image_data) as source_image: # if the image data is not animated then we don't need to generate still if source_image.get_n_pages() == 1: return None if source_image.width == source_image.height and maybe_skip and emoji.is_animated: return None still_frame = pyvips.Image.thumbnail_buffer( image_data, size, height=size, ) if still_frame.width != still_frame.height: if not still_frame.hasalpha(): still_frame = still_frame.addalpha() still_frame = still_frame.gravity( pyvips.CompassDirection.CENTRE, size, size, extend=pyvips.Extend.BACKGROUND, background=[0, 0, 0, 0], ) return still_frame.write_to_buffer(".png") def fix_still_images_local(apps: StateApps) -> None: assert settings.LOCAL_AVATARS_DIR is not None RealmEmoji = apps.get_model("zerver", "RealmEmoji") ids_to_mark_animated: list[int] = [] for total_processed, emoji in enumerate(RealmEmoji.objects.all().iterator()): if total_processed % 100 == 0: print(f"Processed {total_processed} realm emoji") base_path = os.path.join(settings.LOCAL_AVATARS_DIR, str(emoji.realm_id), "emoji/images") original_path = f"{base_path}/{emoji.file_name}.original" try: still_bytes: bytes | None = None maybe_skip = True if not os.path.exists(original_path): maybe_skip = False original_path = f"{base_path}/{emoji.file_name}" if not os.path.exists(original_path): raise SkipImageError("Emoji file does not exist") with open(original_path, "rb") as fh: original_bytes = fh.read() still_bytes = extract_still_from_animated(emoji, original_bytes, maybe_skip) if still_bytes is None: continue if not emoji.is_animated: ids_to_mark_animated.append(emoji.id) still_dir = f"{base_path}/still" os.makedirs(still_dir, exist_ok=True) filename_no_extension = os.path.splitext(emoji.file_name)[0] still_path = f"{still_dir}/{filename_no_extension}.png" with open(still_path, "wb") as fh: fh.write(still_bytes) except SkipImageError as e: logging.warning( "Failed to regenerate still image for emoji id %d (%s): %s", emoji.id, emoji.file_name, e, ) if ids_to_mark_animated: RealmEmoji.objects.filter(id__in=ids_to_mark_animated).update(is_animated=True) def fix_still_images_s3(apps: StateApps) -> None: avatar_bucket = boto3.resource( "s3", aws_access_key_id=settings.S3_KEY, aws_secret_access_key=settings.S3_SECRET_KEY, region_name=settings.S3_REGION, endpoint_url=settings.S3_ENDPOINT_URL, config=Config( signature_version=None, s3={"addressing_style": settings.S3_ADDRESSING_STYLE}, ), ).Bucket(settings.S3_AVATAR_BUCKET) RealmEmoji = apps.get_model("zerver", "RealmEmoji") ids_to_mark_animated: list[int] = [] for total_processed, emoji in enumerate(RealmEmoji.objects.all().iterator()): if total_processed % 100 == 0: print(f"Processed {total_processed} realm emoji") base_path = os.path.join(str(emoji.realm_id), "emoji/images") original_path = f"{base_path}/{emoji.file_name}.original" try: still_bytes: bytes | None = None maybe_skip = True try: obj = avatar_bucket.Object(original_path).get() except botocore.exceptions.ClientError: try: maybe_skip = False original_path = f"{base_path}/{emoji.file_name}" obj = avatar_bucket.Object(original_path).get() except botocore.exceptions.ClientError as e: raise SkipImageError(f"Failed to read original file: {e}") original_bytes = obj["Body"].read() still_bytes = extract_still_from_animated(emoji, original_bytes, maybe_skip) if still_bytes is None: continue if not emoji.is_animated: ids_to_mark_animated.append(emoji.id) filename_no_extension = os.path.splitext(emoji.file_name)[0] still_key = f"{base_path}/still/{filename_no_extension}.png" metadata = obj.get("Metadata", {}) avatar_bucket.Object(still_key).put( Metadata=metadata, ContentType="image/png", CacheControl="public, max-age=31536000, immutable", Body=still_bytes, ) except SkipImageError as e: logging.warning( "Failed to regenerate still image for emoji id %d (%s): %s", emoji.id, emoji.file_name, e, ) if ids_to_mark_animated: RealmEmoji.objects.filter(id__in=ids_to_mark_animated).update(is_animated=True) def fix_still_images(apps: StateApps, schema_editor: BaseDatabaseSchemaEditor) -> None: if settings.LOCAL_AVATARS_DIR is not None: fix_still_images_local(apps) else: fix_still_images_s3(apps) class Migration(migrations.Migration): atomic = False elidable = True dependencies = [ ("zerver", "0791_alter_archivedusermessage_user_profile_and_more"), ] operations = [ migrations.RunPython( fix_still_images, reverse_code=migrations.RunPython.noop, elidable=True ), ]