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):