mirror of
https://github.com/uhIgnacio/EmoteManager.git
synced 2024-08-15 02:23:13 +00:00
add zip/tar archive support 🎉
This commit is contained in:
parent
7dd40ce0e5
commit
10ed6bb63f
5 changed files with 199 additions and 40 deletions
|
@ -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.
|
<code>@Emote Manager add :some_emote:</code> instead.
|
||||||
</p>
|
</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>
|
<p>
|
||||||
<code>@Emote Manager list</code> gives you a list of all emotes on this server.
|
<code>@Emote Manager list</code> gives you a list of all emotes on this server.
|
||||||
</p>
|
</p>
|
||||||
|
|
119
cogs/emote.py
119
cogs/emote.py
|
@ -3,8 +3,10 @@
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
import cgi
|
import cgi
|
||||||
|
import io
|
||||||
import logging
|
import logging
|
||||||
import weakref
|
import weakref
|
||||||
|
import posixpath
|
||||||
import traceback
|
import traceback
|
||||||
import contextlib
|
import contextlib
|
||||||
import urllib.parse
|
import urllib.parse
|
||||||
|
@ -15,6 +17,7 @@ import discord
|
||||||
from discord.ext import commands
|
from discord.ext import commands
|
||||||
|
|
||||||
import utils
|
import utils
|
||||||
|
import utils.archive
|
||||||
import utils.image
|
import utils.image
|
||||||
from utils import errors
|
from utils import errors
|
||||||
from utils.paginator import ListPaginator
|
from utils.paginator import ListPaginator
|
||||||
|
@ -22,6 +25,12 @@ from utils.paginator import ListPaginator
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
class Emotes(commands.Cog):
|
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):
|
def __init__(self, bot):
|
||||||
self.bot = bot
|
self.bot = bot
|
||||||
self.http = aiohttp.ClientSession(loop=self.bot.loop, read_timeout=30, headers={
|
self.http = aiohttp.ClientSession(loop=self.bot.loop, read_timeout=30, headers={
|
||||||
|
@ -122,15 +131,19 @@ class Emotes(commands.Cog):
|
||||||
elif not args:
|
elif not args:
|
||||||
raise commands.BadArgument('Your message had no emotes and no name!')
|
raise commands.BadArgument('Your message had no emotes and no name!')
|
||||||
|
|
||||||
@staticmethod
|
@classmethod
|
||||||
def parse_add_command_attachment(context, args):
|
def parse_add_command_attachment(cls, context, args):
|
||||||
attachment = context.message.attachments[0]
|
attachment = context.message.attachments[0]
|
||||||
# as far as i can tell, this is how discord replaces filenames when you upload an emote image
|
name = cls.format_emote_filename(''.join(args) if args else attachment.filename)
|
||||||
name = ''.join(args) if args else attachment.filename.split('.')[0].replace(' ', '')
|
|
||||||
url = attachment.url
|
url = attachment.url
|
||||||
|
|
||||||
return name, 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'])
|
@commands.command(name='add-from-ec', aliases=['addfromec'])
|
||||||
async def add_from_ec(self, context, name, *names):
|
async def add_from_ec(self, context, name, *names):
|
||||||
"""Copies one or more emotes from Emote Collector to your server.
|
"""Copies one or more emotes from Emote Collector to your server.
|
||||||
|
@ -163,56 +176,100 @@ class Emotes(commands.Cog):
|
||||||
|
|
||||||
await context.send(message)
|
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):
|
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 to add an emote. Returns a string that should be sent to the user."""
|
||||||
try:
|
try:
|
||||||
emote = await self.add_from_url(guild, name, url, author_id, reason=reason)
|
image_data = await self.fetch_safe(url)
|
||||||
except discord.HTTPException as ex:
|
except errors.InvalidFileError:
|
||||||
return (
|
raise errors.InvalidImageError
|
||||||
'An error occurred while creating the emote:\n'
|
|
||||||
+ utils.format_http_exception(ex))
|
if type(image_data) is str: # error case
|
||||||
except errors.ImageResizeTimeoutError:
|
return image_data
|
||||||
raise
|
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:
|
except asyncio.TimeoutError:
|
||||||
return 'Error: retrieving the image took too long.'
|
return 'Error: retrieving the image took too long.'
|
||||||
except ValueError:
|
except ValueError:
|
||||||
return 'Error: Invalid URL.'
|
return 'Error: Invalid URL.'
|
||||||
else:
|
|
||||||
|
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 f'Emote {emote} successfully created.'
|
||||||
|
|
||||||
async def add_from_url(self, guild, name, url, author_id, *, reason=None):
|
async def fetch(self, url, valid_mimetypes=None):
|
||||||
image_data = await self.fetch_emote(url)
|
valid_mimetypes = valid_mimetypes or self.IMAGE_MIMETYPES
|
||||||
emote = await self.create_emote_from_bytes(guild, name, author_id, image_data, reason=reason)
|
|
||||||
|
|
||||||
return emote
|
|
||||||
|
|
||||||
async def fetch_emote(self, url):
|
|
||||||
def validate_headers(response):
|
def validate_headers(response):
|
||||||
if response.reason != 'OK':
|
if response.reason != 'OK':
|
||||||
raise errors.HTTPException(response.status)
|
raise errors.HTTPException(response.status)
|
||||||
# some dumb servers also send '; charset=UTF-8' which we should ignore
|
# some dumb servers also send '; charset=UTF-8' which we should ignore
|
||||||
mimetype, options = cgi.parse_header(response.headers.get('Content-Type', ''))
|
mimetype, options = cgi.parse_header(response.headers.get('Content-Type', ''))
|
||||||
if mimetype not in {'image/png', 'image/jpeg', 'image/gif'}:
|
if mimetype not in valid_mimetypes:
|
||||||
raise errors.InvalidImageError
|
raise errors.InvalidFileError
|
||||||
|
|
||||||
|
async def validate(request):
|
||||||
try:
|
try:
|
||||||
async with self.http.head(url, timeout=5) as response:
|
async with request as response:
|
||||||
validate_headers(response)
|
|
||||||
except aiohttp.ServerDisconnectedError as exception:
|
|
||||||
validate_headers(exception.message)
|
|
||||||
|
|
||||||
async with self.http.get(url) as response:
|
|
||||||
validate_headers(response)
|
validate_headers(response)
|
||||||
return await response.read()
|
return await response.read()
|
||||||
|
except aiohttp.ClientError as exc:
|
||||||
|
raise errors.EmoteManagerError('An error occurred while retrieving the file: {exc}')
|
||||||
|
|
||||||
|
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):
|
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)
|
image_data = await utils.image.resize_in_subprocess(image_data)
|
||||||
if reason is None:
|
if reason is None:
|
||||||
reason = f'Created by {utils.format_user(self.bot, author_id)}'
|
reason = f'Created by {utils.format_user(self.bot, author_id)}'
|
||||||
return await guild.create_custom_emoji(
|
return await guild.create_custom_emoji(name=name, image=image_data, reason=reason)
|
||||||
name=name,
|
|
||||||
image=image_data,
|
|
||||||
reason=reason)
|
|
||||||
|
|
||||||
@commands.command(aliases=('delete', 'delet', 'rm'))
|
@commands.command(aliases=('delete', 'delet', 'rm'))
|
||||||
async def remove(self, context, emote, *emotes):
|
async def remove(self, context, emote, *emotes):
|
||||||
|
|
|
@ -2,6 +2,7 @@
|
||||||
# encoding: utf-8
|
# encoding: utf-8
|
||||||
|
|
||||||
from .misc import *
|
from .misc import *
|
||||||
|
from . import archive
|
||||||
from . import emote
|
from . import emote
|
||||||
from . import errors
|
from . import errors
|
||||||
from . import paginator
|
from . import paginator
|
||||||
|
|
92
utils/archive.py
Normal file
92
utils/archive.py
Normal file
|
@ -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()
|
|
@ -4,7 +4,6 @@ from discord.ext import commands
|
||||||
|
|
||||||
import utils
|
import utils
|
||||||
|
|
||||||
|
|
||||||
class MissingManageEmojisPermission(commands.MissingPermissions):
|
class MissingManageEmojisPermission(commands.MissingPermissions):
|
||||||
"""The invoker or the bot doesn't have permissions to manage server emojis."""
|
"""The invoker or the bot doesn't have permissions to manage server emojis."""
|
||||||
|
|
||||||
|
@ -33,16 +32,21 @@ class EmoteNotFoundError(EmoteManagerError):
|
||||||
def __init__(self, name):
|
def __init__(self, name):
|
||||||
super().__init__(f'An emote called `{name}` does not exist in this server.')
|
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"""
|
"""The image is not a GIF, PNG, or JPG"""
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
super().__init__('The image supplied was not a GIF, PNG, or JPG.')
|
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):
|
class PermissionDeniedError(EmoteManagerError):
|
||||||
"""Raised when a user tries to modify an emote without the Manage Emojis permission"""
|
"""Raised when a user tries to modify an emote without the Manage Emojis permission"""
|
||||||
def __init__(self, name):
|
def __init__(self, name):
|
||||||
|
|
Loading…
Reference in a new issue