From 10ed6bb63fddbdfd6d41fae7df769cd3f0f3f458 Mon Sep 17 00:00:00 2001 From: Benjamin Mintz Date: Sun, 4 Aug 2019 10:15:15 +0000 Subject: [PATCH] =?UTF-8?q?add=20zip/tar=20archive=20support=20?= =?UTF-8?q?=F0=9F=8E=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 5 ++ cogs/emote.py | 123 +++++++++++++++++++++++++++++++++------------- utils/__init__.py | 1 + utils/archive.py | 92 ++++++++++++++++++++++++++++++++++ utils/errors.py | 18 ++++--- 5 files changed, 199 insertions(+), 40 deletions(-) create mode 100644 utils/archive.py diff --git a/README.md b/README.md index 09b821c..20de0c3 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,11 @@ To add the bot to your server, visit https://discordapp.com/oauth2/authorize?cli @Emote Manager add :some_emote: instead.

+

+ To add several emotes from a zip or tar archive, run @Emote Manager add-archive with an attached file. + You can also pass a URL to a zip or tar archive. +

+

@Emote Manager list gives you a list of all emotes on this server.

diff --git a/cogs/emote.py b/cogs/emote.py index efa52be..a5ce10c 100644 --- a/cogs/emote.py +++ b/cogs/emote.py @@ -3,8 +3,10 @@ import asyncio import cgi +import io import logging import weakref +import posixpath import traceback import contextlib import urllib.parse @@ -15,6 +17,7 @@ import discord from discord.ext import commands import utils +import utils.archive import utils.image from utils import errors from utils.paginator import ListPaginator @@ -22,6 +25,12 @@ from utils.paginator import ListPaginator logger = logging.getLogger(__name__) class Emotes(commands.Cog): + IMAGE_MIMETYPES = {'image/png', 'image/jpeg', 'image/gif'} + # TAR_MIMETYPES = {'application/x-tar', 'application/x-xz', 'application/gzip', 'application/x-bzip2'} + TAR_MIMETYPES = {'application/x-tar'} + ZIP_MIMETYPES = {'application/zip', 'application/octet-stream', 'application/x-zip-compressed', 'multipart/x-zip'} + ARCHIVE_MIMETYPES = TAR_MIMETYPES | ZIP_MIMETYPES + def __init__(self, bot): self.bot = bot self.http = aiohttp.ClientSession(loop=self.bot.loop, read_timeout=30, headers={ @@ -122,15 +131,19 @@ class Emotes(commands.Cog): elif not args: raise commands.BadArgument('Your message had no emotes and no name!') - @staticmethod - def parse_add_command_attachment(context, args): + @classmethod + def parse_add_command_attachment(cls, context, args): attachment = context.message.attachments[0] - # as far as i can tell, this is how discord replaces filenames when you upload an emote image - name = ''.join(args) if args else attachment.filename.split('.')[0].replace(' ', '') + name = cls.format_emote_filename(''.join(args) if args else attachment.filename) url = attachment.url return name, url + @staticmethod + def format_emote_filename(filename): + """format a filename to an emote name as discord does when you upload an emote image""" + return posixpath.splitext(filename)[0].replace(' ', '') + @commands.command(name='add-from-ec', aliases=['addfromec']) async def add_from_ec(self, context, name, *names): """Copies one or more emotes from Emote Collector to your server. @@ -163,56 +176,100 @@ class Emotes(commands.Cog): await context.send(message) + @commands.command(name='add-zip', aliases=['add-tar', 'add-from-zip', 'add-from-tar']) + async def add_archive(self, context, url=None): + """Add several emotes from a .zip or .tar archive. + + You may either pass a URL to an archive or upload one as an attachment. + All .gif, .png, and .jpg files in the archive will be uploaded as emotes. + """ + if url and context.message.attachments: + raise commands.BadArgument('Either a URL or an attachment must be given, not both.') + if not url and not context.message.attachments: + raise commands.BadArgument('A URL or attachment must be given.') + + url = url or context.message.attachments[0].url + archive = await self.fetch_safe(url, valid_mimetypes=self.ARCHIVE_MIMETYPES) + if type(archive) is str: # error case + await context.send(archive) + return + + await self.add_from_archive(context, archive) + + async def add_from_archive(self, context, archive): + limit = 50_000_000 # prevent someone from trying to make a giant compressed file + async for name, img, error in utils.archive.extract_async(io.BytesIO(archive), size_limit=limit): + if error is None: + name = self.format_emote_filename(name) + async with context.typing(): + await context.send(await self.add_safe_bytes(context.guild, name, context.author.id, img)) + continue + + if isinstance(error, errors.FileTooBigError): + await context.send( + f'{name}: file too big. ' + f'The limit is {humanize.naturalsize(error.limit)} ' + f'but this file is {humainze.naturalsize(error.size)}.') + continue + + await context.send(f'{name}: {error}') + async def add_safe(self, guild, name, url, author_id, *, reason=None): """Try to add an emote. Returns a string that should be sent to the user.""" try: - emote = await self.add_from_url(guild, name, url, author_id, reason=reason) - except discord.HTTPException as ex: - return ( - 'An error occurred while creating the emote:\n' - + utils.format_http_exception(ex)) - except errors.ImageResizeTimeoutError: - raise + image_data = await self.fetch_safe(url) + except errors.InvalidFileError: + raise errors.InvalidImageError + + if type(image_data) is str: # error case + return image_data + return await self.add_safe_bytes(guild, name, author_id, image_data, reason=reason) + + async def fetch_safe(self, url, valid_mimetypes=None): + """Try to fetch a URL. On error return a string that should be sent to the user.""" + try: + return await self.fetch(url, valid_mimetypes=valid_mimetypes) except asyncio.TimeoutError: return 'Error: retrieving the image took too long.' except ValueError: return 'Error: Invalid URL.' - else: - return f'Emote {emote} successfully created.' - async def add_from_url(self, guild, name, url, author_id, *, reason=None): - image_data = await self.fetch_emote(url) - emote = await self.create_emote_from_bytes(guild, name, author_id, image_data, reason=reason) + async def add_safe_bytes(self, guild, name, author_id, image_data: bytes, *, reason=None): + """Try to add an emote from bytes. On error, return a string that should be sent to the user.""" + try: + emote = await self.create_emote_from_bytes(guild, name, author_id, image_data, reason=reason) + except discord.HTTPException as ex: + return ( + 'An error occurred while creating the emote:\n' + + utils.format_http_exception(ex)) + return f'Emote {emote} successfully created.' - return emote - - async def fetch_emote(self, url): + async def fetch(self, url, valid_mimetypes=None): + valid_mimetypes = valid_mimetypes or self.IMAGE_MIMETYPES def validate_headers(response): if response.reason != 'OK': raise errors.HTTPException(response.status) # 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 + if mimetype not in valid_mimetypes: + raise errors.InvalidFileError - 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 validate(request): + try: + async with request as response: + validate_headers(response) + return await response.read() + except aiohttp.ClientError as exc: + raise errors.EmoteManagerError('An error occurred while retrieving the file: {exc}') - async with self.http.get(url) as response: - validate_headers(response) - return await response.read() + await validate(self.http.head(url, timeout=5)) + return await validate(self.http.get(url)) 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, - reason=reason) + return await guild.create_custom_emoji(name=name, image=image_data, reason=reason) @commands.command(aliases=('delete', 'delet', 'rm')) async def remove(self, context, emote, *emotes): diff --git a/utils/__init__.py b/utils/__init__.py index d140b81..2e0a62c 100644 --- a/utils/__init__.py +++ b/utils/__init__.py @@ -2,6 +2,7 @@ # encoding: utf-8 from .misc import * +from . import archive from . import emote from . import errors from . import paginator diff --git a/utils/archive.py b/utils/archive.py new file mode 100644 index 0000000..f142068 --- /dev/null +++ b/utils/archive.py @@ -0,0 +1,92 @@ +#!/usr/bin/env python3 + +import asyncio +import collections +import enum +import posixpath +import tarfile +import typing.io +import zipfile +from typing import Iterable, Tuple, Optional + +from . import errors + +ArchiveInfo = collections.namedtuple('ArchiveInfo', 'filename content error') + +def extract(archive: typing.io.BinaryIO, *, size_limit=None) \ + -> Iterable[Tuple[str, Optional[bytes], Optional[BaseException]]]: + """ + extract a binary file-like object representing a zip or uncompressed tar archive, yielding filenames and contents. + + yields ArchiveInfo objects: (filename: str, content: typing.Optional[bytes], error: ) + if size_limit is not None and the size limit is exceeded, or for any other error, yield None for content + on success, error will be None + """ + + try: + yield from extract_zip(archive, size_limit=size_limit) + return + except zipfile.BadZipFile: + pass + finally: + archive.seek(0) + + try: + yield from extract_tar(archive, size_limit=size_limit) + except tarfile.ReadError as exc: + raise ValueError('not a valid zip or tar file') from exc + finally: + archive.seek(0) + +def extract_zip(archive, *, size_limit=None): + with zipfile.ZipFile(archive) as zip: + members = [m for m in zip.infolist() if not m.is_dir()] + for member in members: + if size_limit is not None and member.file_size >= size_limit: + yield ArchiveInfo( + filename=member.filename, + content=None, + error=errors.FileTooBigError(member.file_size, size_limit)) + continue + + try: + content = zip.open(member).read() + except RuntimeError as exc: # why no specific exceptions smh + yield ArchiveInfo(filename=member.filename, content=None, error=exc) + else: # this else is required to avoid UnboundLocalError for some reason + yield ArchiveInfo(filename=member.filename, content=content, error=None) + +def extract_tar(archive, *, size_limit=None): + with tarfile.open(fileobj=archive) as tar: + members = [f for f in tar.getmembers() if f.isfile()] + for member in members: + if size_limit is not None and member.size >= size_limit: + yield ArchiveInfo( + filename=member.name, + content=None, + error=errors.FileTooBigError(member.size, size_limit)) + continue + + yield ArchiveInfo(member.name, content=tar.extractfile(member).read(), error=None) + +async def extract_async(archive: typing.io.BinaryIO, size_limit=None): + gen = await asyncio.get_event_loop().run_in_executor(None, extract, archive) + for x in gen: + yield await asyncio.sleep(0, x) + +def main(): + import io + import sys + + import humanize + + arc = io.BytesIO(sys.stdin.detach().read()) + for name, data, error in extract(arc): + if error is not None: + print(f'{name}: {error}') + continue + + print(f'{name}: {humanize.naturalsize(len(data)):>10}') + +if __name__ == '__main__': + main() diff --git a/utils/errors.py b/utils/errors.py index d43aaf6..219da96 100644 --- a/utils/errors.py +++ b/utils/errors.py @@ -4,7 +4,6 @@ from discord.ext import commands import utils - class MissingManageEmojisPermission(commands.MissingPermissions): """The invoker or the bot doesn't have permissions to manage server emojis.""" @@ -33,16 +32,21 @@ class EmoteNotFoundError(EmoteManagerError): def __init__(self, name): super().__init__(f'An emote called `{name}` does not exist in this server.') -class InvalidImageError(EmoteManagerError): +class FileTooBigError(EmoteManagerError): + def __init__(self, size, limit): + self.size = size + self.limit = limit + +class InvalidFileError(EmoteManagerError): + """The file is not a zip, tar, GIF, PNG, or JPG file.""" + def __init__(self): + super().__init__('Invalid file given.') + +class InvalidImageError(InvalidFileError): """The image is not a GIF, PNG, or JPG""" def __init__(self): super().__init__('The image supplied was not a GIF, PNG, or JPG.') -class NoMoreSlotsError(EmoteManagerError): - """Raised in case all slots of a particular type (static/animated) are full""" - def __init__(self): - super().__init__('No more backend slots available.') - class PermissionDeniedError(EmoteManagerError): """Raised when a user tries to modify an emote without the Manage Emojis permission""" def __init__(self, name):