diff --git a/src/main/cogs/booru.py b/src/main/cogs/booru.py index a82564a..4442497 100644 --- a/src/main/cogs/booru.py +++ b/src/main/cogs/booru.py @@ -1,12 +1,12 @@ import asyncio import json import re +import sys import traceback as tb from contextlib import suppress -from fractions import gcd from datetime import datetime as dt from datetime import timedelta as td -import sys +from fractions import gcd import discord as d from discord import errors as err @@ -41,11 +41,11 @@ class MsG: 'cogs/blacklists.pkl', {'global_blacklist': set(), 'guild_blacklist': {}, 'user_blacklist': {}}) self.aliases = u.setdefault('cogs/aliases.pkl', {}) - if u.tasks['auto_qual']: - for channel in u.tasks['auto_qual']: + if u.tasks['auto_rev']: + for channel in u.tasks['auto_rev']: temp = self.bot.get_channel(channel) self.bot.loop.create_task(self.queue_for_reversification(temp)) - print('AUTO-QUALITIFYING : #{}'.format(temp.name)) + print('AUTO-REVERSIFYING : #{}'.format(temp.name)) self.reversifying = True self.bot.loop.create_task(self._reversify()) # if not self.updating: @@ -219,7 +219,8 @@ class MsG: try: await dest.trigger_typing() - ident = ident if not ident.isdigit() else re.search('show/([0-9]+)', ident).group(1) + ident = ident if not ident.isdigit() else re.search( + 'show/([0-9]+)', ident).group(1) post = await u.fetch('https://e621.net/post/show.json', params={'id': ident}, json=True) embed = d.Embed( @@ -227,7 +228,7 @@ class MsG: embed.set_thumbnail(url=post['file_url']) embed.set_author(name=f'{u.get_aspectratio(post["width"], post["height"])} \N{ZERO WIDTH SPACE} {post["width"]} x {post["height"]}', url=f'https://e621.net/post?tags=ratio:{post["width"]/post["height"]:.2f}', icon_url=ctx.author.avatar_url) - embed.set_footer(text=str(post['score']), + embed.set_footer(text=post['score'], icon_url=self._get_score(post['score'])) # except @@ -342,7 +343,7 @@ class MsG: embed.set_image(url=post['file_url']) embed.set_author(name=f'{u.get_aspectratio(post["width"], post["height"])} \N{ZERO WIDTH SPACE} {post["width"]} x {post["height"]}', url=f'https://e621.net/post?tags=ratio:{post["width"]/post["height"]:.2f}', icon_url=ctx.author.avatar_url) - embed.set_footer(text=str(post['score']), + embed.set_footer(text=post['score'], icon_url=self._get_score(post['score'])) await dest.send('**Probable match**', embed=embed) @@ -352,14 +353,14 @@ class MsG: except exc.MatchError as e: await ctx.send('**No probable match for:** `{}`'.format(e), delete_after=7) - finally: - await asyncio.sleep(self.RATE_LIMIT) - if not c: await ctx.message.add_reaction('\N{CROSS MARK}') except exc.MissingArgument: - await ctx.send('**Invalid url or file** Be sure the link directs to an image file', delete_after=7) + await ctx.send('**Invalid url or file.** Be sure the link directs to an image file', delete_after=7) + await ctx.message.add_reaction('\N{CROSS MARK}') + except exc.SizeError as e: + await ctx.send(f'`{e}` **too large.** Maximum is 8 MB', delete_after=7) await ctx.message.add_reaction('\N{CROSS MARK}') @commands.command(name='reversify', aliases=['revify', 'risify', 'rify']) @@ -399,7 +400,17 @@ class MsG: try: await dest.trigger_typing() - await dest.send('`{} / {}` **Probable match from** {}\n{}'.format(n, len(links), message.author.display_name, await scraper.get_post(url))) + post = await scraper.get_post(url) + + embed = d.Embed( + title=', '.join(post['artist']), url=f'https://e621.net/post/show/{post["id"]}', color=ctx.me.color if isinstance(ctx.channel, d.TextChannel) else u.color) + embed.set_image(url=post['file_url']) + embed.set_author(name=f'{u.get_aspectratio(post["width"], post["height"])} \N{ZERO WIDTH SPACE} {post["width"]} x {post["height"]}', + url=f'https://e621.net/post?tags=ratio:{post["width"]/post["height"]:.2f}', icon_url=ctx.author.avatar_url) + embed.set_footer( + text=post['score'], icon_url=self._get_score(post['score'])) + + await dest.send(f'**Probable match from** {message.author.display_name}', embed=embed) await message.add_reaction('\N{WHITE HEAVY CHECK MARK}') if remove: @@ -410,10 +421,13 @@ class MsG: await ctx.send('`{} / {}` **No probable match for:** `{}`'.format(n, len(links), e), delete_after=7) await message.add_reaction('\N{CROSS MARK}') c -= 1 + except exc.SizeError as e: + await ctx.send(f'`{e}` **too large.** Maximum is 8 MB', delete_after=7) + await message.add_reaction('\N{CROSS MARK}') + c -= 1 finally: n += 1 - await asyncio.sleep(self.RATE_LIMIT) if c <= 0: await ctx.message.add_reaction('\N{CROSS MARK}') @@ -444,20 +458,31 @@ class MsG: post = await scraper.get_post(url) - await message.channel.send('**Probable match from** {}\n{}'.format(message.author.display_name, await scraper.get_image(post))) - with suppress(err.NotFound): - await message.add_reaction('\N{WHITE HEAVY CHECK MARK}') + embed = d.Embed( + title=', '.join(post['artist']), url=f'https://e621.net/post/show/{post["id"]}', color=message.channel.guild.me.color if isinstance(message.channel, d.TextChannel) else u.color) + embed.set_image(url=post['file_url']) + embed.set_author(name=f'{u.get_aspectratio(post["width"], post["height"])} \N{ZERO WIDTH SPACE} {post["width"]} x {post["height"]}', + url=f'https://e621.net/post?tags=ratio:{post["width"]/post["height"]:.2f}', icon_url=message.author.avatar_url) + embed.set_footer(text=post['score'], + icon_url=self._get_score(post['score']['score'])) - with suppress(err.NotFound): - await message.delete() + await message.channel.send('**Probable match from** {}'.format(message.author.display_name), embed=embed) + + await message.add_reaction('\N{WHITE HEAVY CHECK MARK}') except exc.MatchError as e: await message.channel.send('**No probable match for:** `{}`'.format(e), delete_after=7) await message.add_reaction('\N{CROSS MARK}') + except exc.SizeError as e: + await message.channel.send(f'`{e}` **too large.** Maximum is 8 MB', delete_after=7) + await message.add_reaction('\N{CROSS MARK}') finally: await asyncio.sleep(self.RATE_LIMIT) + with suppress(err.NotFound): + await message.delete() + print('STOPPED : reversifying') async def queue_for_reversification(self, channel): @@ -475,18 +500,18 @@ class MsG: await message.add_reaction('\N{HOURGLASS WITH FLOWING SAND}') except exc.Abort: - u.tasks['auto_qual'].remove(channel.id) + u.tasks['auto_rev'].remove(channel.id) u.dump(u.tasks, 'cogs/tasks.pkl') - if not u.tasks['auto_qual']: + if not u.tasks['auto_rev']: self.reversifying = False print('STOPPED : reversifying #{}'.format(channel.name)) await channel.send('**Stopped queueing messages for reversification in** {}'.format(channel.mention), delete_after=5) - @commands.command(name='autoreversify', aliases=['autoqual']) + @commands.command(name='autoreversify', aliases=['autorev']) @commands.has_permissions(manage_channels=True) async def auto_reversify(self, ctx): - if ctx.channel.id not in u.tasks['auto_qual']: - u.tasks['auto_qual'].append(ctx.channel.id) + if ctx.channel.id not in u.tasks['auto_rev']: + u.tasks['auto_rev'].append(ctx.channel.id) u.dump(u.tasks, 'cogs/tasks.pkl') self.bot.loop.create_task( self.queue_for_reversification(ctx.channel)) diff --git a/src/main/misc/exceptions.py b/src/main/misc/exceptions.py index dcc8c84..7959441 100644 --- a/src/main/misc/exceptions.py +++ b/src/main/misc/exceptions.py @@ -10,6 +10,9 @@ async def send_error(ctx, error): # class NSFW(errext.CheckFailure): # pass +class SizeError(Exception): + pass + class Wrong(Exception): pass diff --git a/src/main/run.py b/src/main/run.py index 898018d..93ccc58 100644 --- a/src/main/run.py +++ b/src/main/run.py @@ -114,6 +114,17 @@ async def on_command_completion(ctx): if ctx.command.name != 'lastcommand': u.last_commands[ctx.author.id] = ctx +@bot.event +async def on_guild_remove(guild): + print(f'LEFT : {guild.name}') + + for task, idents in u.tasks.items(): + for channel in guild.channels: + if channel.id in idents: + idents.remove(channel.id) + print(f'STOPPED : {task} in #{channel.id}') + u.dump(u.tasks, 'cogs/tasks.pkl') + # d.opus.load_opus('opus') diff --git a/src/main/utils/scraper.py b/src/main/utils/scraper.py index 0482e86..ccb60fe 100644 --- a/src/main/utils/scraper.py +++ b/src/main/utils/scraper.py @@ -3,12 +3,21 @@ import re from bs4 import BeautifulSoup from lxml import html +from hurry.filesize import size, alternative from misc import exceptions as exc from utils import utils as u async def get_post(url): + try: + image = await u.fetch(url, response=True) + filesize = int(image.headers['Content-Length']) + if filesize > 8192 * 1024: + raise exc.SizeError(size(filesize, system=alternative)) + except ValueError: + raise exc.MissingArgument + await asyncio.sleep(u.RATE_LIMIT) content = await u.fetch('http://iqdb.harry.lu', params={'url': url}) diff --git a/src/main/utils/utils.py b/src/main/utils/utils.py index d8c7c91..83e34d7 100644 --- a/src/main/utils/utils.py +++ b/src/main/utils/utils.py @@ -67,7 +67,7 @@ def dump(obj, filename, *, json=False): settings = setdefault('settings.pkl', {'del_ctx': [], 'prefixes': {}}) -tasks = setdefault('cogs/tasks.pkl', {'auto_del': [], 'auto_qual': [], 'auto_rev': []}) +tasks = setdefault('cogs/tasks.pkl', {'auto_del': [], 'auto_rev': []}) temp = setdefault('temp.pkl', {}) RATE_LIMIT = 2.2 @@ -107,9 +107,11 @@ def close(loop): print('Finished cancelling tasks.') -async def fetch(url, *, params={}, json=False): +async def fetch(url, *, params={}, json=False, response=False): async with session.get(url, params=params, headers={'User-Agent': 'Myned/Modumind/dev'}) as r: - if json: + if response: + return r + elif json: return await r.json() return await r.read()