diff --git a/cogs/emote.py b/cogs/emote.py index 0d5eddd..37edb31 100644 --- a/cogs/emote.py +++ b/cogs/emote.py @@ -1,9 +1,8 @@ #!/usr/bin/env python3 # encoding: utf-8 -import io -import imghdr import asyncio +import cgi import logging import weakref import traceback @@ -186,29 +185,31 @@ class Emotes(commands.Cog): return emote async def fetch_emote(self, url): - # credits to @Liara#0001 (ID 136900814408122368) for most of this part - # https://gitlab.com/Pandentia/element-zero/blob/47bc8eeeecc7d353ec66e1ef5235adab98ca9635/element_zero/cogs/emoji.py#L217-228 - async with self.http.head(url, timeout=5) as response: + def validate_headers(response): if response.reason != 'OK': raise errors.HTTPException(response.status) - if response.headers.get('Content-Type') not in ('image/png', 'image/jpeg', 'image/gif'): + # some dumb servers also send '; charset=UTF-8' which we should ignore + mimetype, options = cgi.parse_header(response.headers.get('Content-Type', '')) + if mimetype not in {'image/png', 'image/jpeg', 'image/gif'}: raise errors.InvalidImageError - async with self.http.get(url) as response: - if response.reason != 'OK': - raise errors.HTTPException(response.status) - return io.BytesIO(await response.read()) + try: + async with self.http.head(url, timeout=5) as response: + validate_headers(response) + except aiohttp.ServerDisconnectedError as exception: + validate_headers(exception.message) - async def create_emote_from_bytes(self, guild, name, author_id, image_data: io.BytesIO, *, reason=None): - # resize_until_small is normally blocking, because wand is. - # run_in_executor is magic that makes it non blocking somehow. - # also, None as the executor arg means "use the loop's default executor" - image_data = await self.bot.loop.run_in_executor(None, utils.image.resize_until_small, image_data) + async with self.http.get(url) as response: + validate_headers(response) + return await response.read() + + async def create_emote_from_bytes(self, guild, name, author_id, image_data: bytes, *, reason=None): + image_data = await utils.image.resize_in_subprocess(image_data) if reason is None: reason = f'Created by {utils.format_user(self.bot, author_id)}' return await guild.create_custom_emoji( name=name, - image=image_data.read(), + image=image_data, reason=reason) @commands.command(aliases=('delete', 'delet', 'rm')) diff --git a/utils/image.py b/utils/image.py old mode 100644 new mode 100755 index 022c167..7d8b546 --- a/utils/image.py +++ b/utils/image.py @@ -1,18 +1,27 @@ +#!/usr/bin/env python3 +# encoding: utf-8 + +import asyncio +import base64 +import contextlib import io import logging -import contextlib +import sys +import typing logger = logging.getLogger(__name__) try: - from wand.image import Image -except ImportError: + import wand.image +except (ImportError, OSError): logger.warn('Failed to import wand.image. Image manipulation functions will be unavailable.') +else: + import wand.exceptions from utils import errors -def resize_until_small(image_data: io.BytesIO): - """If the image_data is bigger than 256KB, resize it until it's not""" +def resize_until_small(image_data: io.BytesIO) -> None: + """If the image_data is bigger than 256KB, resize it until it's not.""" # It's important that we only attempt to resize the image when we have to, # ie when it exceeds the Discord limit of 256KiB. # Apparently some <256KiB images become larger when we attempt to resize them, @@ -21,43 +30,108 @@ def resize_until_small(image_data: io.BytesIO): image_size = size(image_data) while image_size > 256 * 2**10 and max_resolution >= 32: # don't resize past 32x32 or 256KiB logger.debug('image size too big (%s bytes)', image_size) - logger.debug('attempting resize to %s*%s pixels', max_resolution, max_resolution) - image_data = thumbnail(image_data, (max_resolution, max_resolution)) + logger.debug('attempting resize to at most%s*%s pixels', max_resolution, max_resolution) + + try: + thumbnail(image_data, (max_resolution, max_resolution)) + except wand.exceptions.CoderError: + raise errors.InvalidImageError + image_size = size(image_data) max_resolution //= 2 - return image_data -def size(data: io.BytesIO): - """return the size, in bytes, of the data a BytesIO object represents""" - with preserve_position(data): - data.seek(0, io.SEEK_END) - return data.tell() - -def thumbnail(image_data: io.BytesIO, max_size=(128, 128)): +def thumbnail(image_data: io.BytesIO, max_size=(128, 128)) -> None: """Resize an image in place to no more than max_size pixels, preserving aspect ratio.""" - # Credit to @Liara#0001 (ID 136900814408122368) - # https://gitlab.com/Pandentia/element-zero/blob/47bc8eeeecc7d353ec66e1ef5235adab98ca9635/element_zero/cogs/emoji.py#L243-247 - image = Image(blob=image_data) - image.resize(*scale_resolution((image.width, image.height), max_size)) - # we create a new buffer here because there's wand errors otherwise. - # specific error: - # MissingDelegateError: no decode delegate for this image format `' @ error/blob.c/BlobToImage/353 - out = io.BytesIO() - image.save(file=out) - out.seek(0) - return out + with wand.image.Image(blob=image_data) as image: + new_resolution = scale_resolution((image.width, image.height), max_size) + image.resize(*new_resolution) + image_data.truncate(0) + image_data.seek(0) + image.save(file=image_data) + + # allow resizing the original image more than once for memory profiling + image_data.seek(0) def scale_resolution(old_res, new_res): - # https://stackoverflow.com/a/6565988 """Resize a resolution, preserving aspect ratio. Returned w,h will be <= new_res""" + # https://stackoverflow.com/a/6565988 + old_width, old_height = old_res new_width, new_height = new_res + old_ratio = old_width / old_height new_ratio = new_width / new_height if new_ratio > old_ratio: return (old_width * new_height//old_height, new_height) return new_width, old_height * new_width//old_width +def mime_type_for_image(data): + if data.startswith(b'\x89PNG\r\n\x1a\n'): + return 'image/png' + elif data.startswith(b'\xFF\xD8') and data.rstrip(b'\0').endswith(b'\xFF\xD9'): + return 'image/jpeg' + elif data.startswith((b'GIF87a', b'GIF89a')): + return 'image/gif' + elif data.startswith(b'RIFF') and data[8:12] == b'WEBP': + return 'image/webp' + else: + raise errors.InvalidImageError + +def image_to_base64_url(data): + fmt = 'data:{mime};base64,{data}' + mime = mime_type_for_image(data) + b64 = base64.b64encode(data).decode('ascii') + return fmt.format(mime=mime, data=b64) + +def main() -> typing.NoReturn: + """resize an image from stdin and write the resized version to stdout.""" + import sys + + data = io.BytesIO(sys.stdin.buffer.read()) + try: + resize_until_small(data) + except errors.InvalidImageError: + # 2 is used because 1 is already used by python's default error handler + sys.exit(2) + + stdout_write = sys.stdout.buffer.write # getattr optimization + + while True: + buf = data.read(16 * 1024) + if not buf: + break + + stdout_write(buf) + + sys.exit(0) + +async def resize_in_subprocess(image_data: bytes): + proc = await asyncio.create_subprocess_exec( + sys.executable, '-m', __name__, + + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE) + + try: + image_data, err = await asyncio.wait_for(proc.communicate(image_data), timeout=30) + except asyncio.TimeoutError: + proc.kill() + raise errors.ImageResizeTimeoutError + else: + if proc.returncode == 2: + raise errors.InvalidImageError + if proc.returncode != 0: + raise RuntimeError(err.decode('utf-8') + f'Return code: {proc.returncode}') + + return image_data + +def size(fp): + """return the size, in bytes, of the data a file-like object represents""" + with preserve_position(fp): + fp.seek(0, io.SEEK_END) + return fp.tell() + class preserve_position(contextlib.AbstractContextManager): def __init__(self, fp): self.fp = fp @@ -65,3 +139,6 @@ class preserve_position(contextlib.AbstractContextManager): def __exit__(self, *excinfo): self.fp.seek(self.old_pos) + +if __name__ == '__main__': + main()