2020-05-12 23:55:08 +00:00
|
|
|
|
# © 2018–2020 io mintz <io@mintz.cc>
|
|
|
|
|
#
|
|
|
|
|
# Emote Manager is free software: you can redistribute it and/or modify
|
|
|
|
|
# it under the terms of the GNU Affero General Public License as
|
|
|
|
|
# published by the Free Software Foundation, either version 3 of the
|
|
|
|
|
# License, or (at your option) any later version.
|
|
|
|
|
#
|
|
|
|
|
# Emote Manager is distributed in the hope that it will be useful,
|
|
|
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
|
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
|
|
|
# GNU Affero General Public License for more details.
|
|
|
|
|
#
|
|
|
|
|
# You should have received a copy of the GNU Affero General Public License
|
|
|
|
|
# along with Emote Manager. If not, see <https://www.gnu.org/licenses/>.
|
2019-06-04 03:08:33 +00:00
|
|
|
|
|
|
|
|
|
import asyncio
|
|
|
|
|
import base64
|
|
|
|
|
import contextlib
|
2019-10-10 00:24:12 +00:00
|
|
|
|
import functools
|
2018-07-30 05:26:15 +00:00
|
|
|
|
import io
|
|
|
|
|
import logging
|
2020-04-29 23:55:31 +00:00
|
|
|
|
import signal
|
2019-06-04 03:08:33 +00:00
|
|
|
|
import sys
|
|
|
|
|
import typing
|
2018-07-30 05:26:15 +00:00
|
|
|
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
|
|
|
|
try:
|
2019-06-04 03:08:33 +00:00
|
|
|
|
import wand.image
|
|
|
|
|
except (ImportError, OSError):
|
2018-07-30 05:26:15 +00:00
|
|
|
|
logger.warn('Failed to import wand.image. Image manipulation functions will be unavailable.')
|
2019-06-04 03:08:33 +00:00
|
|
|
|
else:
|
|
|
|
|
import wand.exceptions
|
2018-07-30 05:26:15 +00:00
|
|
|
|
|
|
|
|
|
from utils import errors
|
|
|
|
|
|
2019-06-04 03:08:33 +00:00
|
|
|
|
def resize_until_small(image_data: io.BytesIO) -> None:
|
|
|
|
|
"""If the image_data is bigger than 256KB, resize it until it's not."""
|
2018-07-30 05:26:15 +00:00
|
|
|
|
# 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,
|
|
|
|
|
# so resizing sometimes does more harm than good.
|
|
|
|
|
max_resolution = 128 # pixels
|
|
|
|
|
image_size = size(image_data)
|
2019-10-10 00:24:12 +00:00
|
|
|
|
if image_size <= 256 * 2**10:
|
|
|
|
|
return
|
2019-06-04 03:08:33 +00:00
|
|
|
|
|
2019-10-10 00:24:12 +00:00
|
|
|
|
try:
|
|
|
|
|
with wand.image.Image(blob=image_data) as original_image:
|
|
|
|
|
while True:
|
|
|
|
|
logger.debug('image size too big (%s bytes)', image_size)
|
|
|
|
|
logger.debug('attempting resize to at most%s*%s pixels', max_resolution, max_resolution)
|
|
|
|
|
|
|
|
|
|
with original_image.clone() as resized:
|
|
|
|
|
resized.transform(resize=f'{max_resolution}x{max_resolution}')
|
|
|
|
|
image_size = len(resized.make_blob())
|
|
|
|
|
if image_size <= 256 * 2**10 or max_resolution < 32: # don't resize past 256KiB or 32×32
|
|
|
|
|
image_data.truncate(0)
|
|
|
|
|
image_data.seek(0)
|
|
|
|
|
resized.save(file=image_data)
|
|
|
|
|
image_data.seek(0)
|
|
|
|
|
break
|
|
|
|
|
|
|
|
|
|
max_resolution //= 2
|
|
|
|
|
except wand.exceptions.CoderError:
|
|
|
|
|
raise errors.InvalidImageError
|
|
|
|
|
|
|
|
|
|
def convert_to_gif(image_data: io.BytesIO) -> None:
|
|
|
|
|
try:
|
|
|
|
|
with wand.image.Image(blob=image_data) as orig, orig.convert('gif') as converted:
|
|
|
|
|
# discord tries to stop us from abusing animated gif slots by detecting single frame gifs
|
|
|
|
|
# so make it two frames
|
|
|
|
|
converted.sequence[0].delay = 0 # show the first frame forever
|
|
|
|
|
converted.sequence.append(wand.image.Image(width=1, height=1))
|
|
|
|
|
|
|
|
|
|
image_data.truncate(0)
|
|
|
|
|
image_data.seek(0)
|
|
|
|
|
converted.save(file=image_data)
|
|
|
|
|
image_data.seek(0)
|
|
|
|
|
except wand.exceptions.CoderError:
|
|
|
|
|
raise errors.InvalidImageError
|
2018-07-30 05:26:15 +00:00
|
|
|
|
|
2019-06-04 03:08:33 +00:00
|
|
|
|
def mime_type_for_image(data):
|
|
|
|
|
if data.startswith(b'\x89PNG\r\n\x1a\n'):
|
|
|
|
|
return 'image/png'
|
2019-08-04 10:38:29 +00:00
|
|
|
|
if data.startswith(b'\xFF\xD8') and data.rstrip(b'\0').endswith(b'\xFF\xD9'):
|
2019-06-04 03:08:33 +00:00
|
|
|
|
return 'image/jpeg'
|
2019-08-04 10:38:29 +00:00
|
|
|
|
if data.startswith((b'GIF87a', b'GIF89a')):
|
2019-06-04 03:08:33 +00:00
|
|
|
|
return 'image/gif'
|
2019-08-04 10:38:29 +00:00
|
|
|
|
raise errors.InvalidImageError
|
2019-06-04 03:08:33 +00:00
|
|
|
|
|
|
|
|
|
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:
|
2019-10-10 00:24:12 +00:00
|
|
|
|
"""resize or convert an image from stdin and write the resized or converted version to stdout."""
|
2019-06-04 03:08:33 +00:00
|
|
|
|
import sys
|
|
|
|
|
|
2019-10-10 00:24:12 +00:00
|
|
|
|
if sys.argv[1] == 'resize':
|
|
|
|
|
f = resize_until_small
|
|
|
|
|
elif sys.argv[1] == 'convert':
|
|
|
|
|
f = convert_to_gif
|
|
|
|
|
else:
|
|
|
|
|
sys.exit(1)
|
|
|
|
|
|
2019-06-04 03:08:33 +00:00
|
|
|
|
data = io.BytesIO(sys.stdin.buffer.read())
|
|
|
|
|
try:
|
2019-10-10 00:24:12 +00:00
|
|
|
|
f(data)
|
2019-06-04 03:08:33 +00:00
|
|
|
|
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)
|
|
|
|
|
|
2019-10-10 00:24:12 +00:00
|
|
|
|
async def process_image_in_subprocess(command_name, image_data: bytes):
|
2019-06-04 03:08:33 +00:00
|
|
|
|
proc = await asyncio.create_subprocess_exec(
|
2019-10-10 00:24:12 +00:00
|
|
|
|
sys.executable, '-m', __name__, command_name,
|
2019-06-04 03:08:33 +00:00
|
|
|
|
|
|
|
|
|
stdin=asyncio.subprocess.PIPE,
|
|
|
|
|
stdout=asyncio.subprocess.PIPE,
|
|
|
|
|
stderr=asyncio.subprocess.PIPE)
|
|
|
|
|
|
|
|
|
|
try:
|
2019-10-10 00:24:12 +00:00
|
|
|
|
image_data, err = await asyncio.wait_for(proc.communicate(image_data), timeout=float('inf'))
|
2019-06-04 03:08:33 +00:00
|
|
|
|
except asyncio.TimeoutError:
|
2020-04-29 23:55:31 +00:00
|
|
|
|
proc.send_signal(signal.SIGINT)
|
2019-10-10 00:24:12 +00:00
|
|
|
|
raise errors.ImageResizeTimeoutError if command_name == 'resize' else errors.ImageConversionTimeoutError
|
2019-06-04 03:08:33 +00:00
|
|
|
|
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
|
|
|
|
|
|
2019-10-10 00:24:12 +00:00
|
|
|
|
resize_in_subprocess = functools.partial(process_image_in_subprocess, 'resize')
|
|
|
|
|
convert_to_gif_in_subprocess = functools.partial(process_image_in_subprocess, 'convert')
|
|
|
|
|
|
2019-06-04 03:08:33 +00:00
|
|
|
|
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()
|
|
|
|
|
|
2018-07-30 05:26:15 +00:00
|
|
|
|
class preserve_position(contextlib.AbstractContextManager):
|
|
|
|
|
def __init__(self, fp):
|
|
|
|
|
self.fp = fp
|
|
|
|
|
self.old_pos = fp.tell()
|
|
|
|
|
|
|
|
|
|
def __exit__(self, *excinfo):
|
|
|
|
|
self.fp.seek(self.old_pos)
|
2019-06-04 03:08:33 +00:00
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
|
main()
|