add zip/tar archive support 🎉

README.md

@@ -21,6 +21,11 @@ To add the bot to your server, visit https://discordapp.com/oauth2/authorize?cli
 	<code>@Emote Manager add :some_emote:</code> instead.
 </p>
 
+<p>
+	To add several emotes from a zip or tar archive, run <code>@Emote Manager add-archive</code> with an attached file.
+	You can also pass a URL to a zip or tar archive.
+</p>
+
 <p>
 	<code>@Emote Manager list</code> gives you a list of all emotes on this server.
 </p>

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

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

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

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