diff --git a/src/main/cogs/booru.py b/src/main/cogs/booru.py index 1661cb7..8b400b8 100644 --- a/src/main/cogs/booru.py +++ b/src/main/cogs/booru.py @@ -23,12 +23,22 @@ class MsG: self.bot = bot self.LIMIT = 100 self.RATE_LIMIT = 2.1 + self.queue = asyncio.Queue() + self.reversing = False self.favorites = u.setdefault('cogs/favorites.pkl', {'tags': set(), 'posts': set()}) self.blacklists = u.setdefault( 'cogs/blacklists.pkl', {'global_blacklist': set(), 'guild_blacklist': {}, 'user_blacklist': {}}) self.aliases = u.setdefault('cogs/aliases.pkl', {}) + 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_on_message(temp)) + print('LOOPING : #{}'.format(temp.name)) + self.bot.loop.create_task(self.reverse()) + self.reversing = True + # Tag search @commands.command(aliases=['rel'], brief='e621 Related tag search', description='e621 | NSFW\nReturn a link search for given tags') @checks.del_ctx() @@ -104,24 +114,22 @@ class MsG: await ctx.send('✅ **Probable match:**\n{}'.format(await scraper.get_post(url))) + await asyncio.sleep(self.RATE_LIMIT) + except exc.MatchError as e: await ctx.send('❌ **No probable match for:** `{}`'.format(e), delete_after=10) - finally: - await asyncio.sleep(self.RATE_LIMIT) - for attachment in ctx.message.attachments: try: await ctx.trigger_typing() await ctx.send('✅ **Probable match:**\n{}'.format(await scraper.get_post(attachment.url))) + await asyncio.sleep(self.RATE_LIMIT) + except exc.MatchError as e: await ctx.send('❌ **No probable match for:** `{}`'.format(e), delete_after=10) - finally: - await asyncio.sleep(self.RATE_LIMIT) - except exc.MissingArgument: await ctx.send('❌ **Invalid url or file.**', delete_after=10) @@ -155,6 +163,8 @@ class MsG: await ctx.send('✅ **Probable match from {}:**\n{}'.format(message.author.display_name, await scraper.get_post(match.group(0)))) + await asyncio.sleep(self.RATE_LIMIT) + if delete: with suppress(err.NotFound): await message.delete() @@ -162,9 +172,6 @@ class MsG: except exc.MatchError as e: await ctx.send('❌ **No probable match for:** `{}`'.format(e), delete_after=10) - finally: - await asyncio.sleep(self.RATE_LIMIT) - for message in attachments: for attachment in message.attachments: try: @@ -172,6 +179,8 @@ class MsG: await ctx.send('✅ **Probable match from {}:**\n{}'.format(message.author.display_name, await scraper.get_post(attachment.url))) + await asyncio.sleep(self.RATE_LIMIT) + if delete: with suppress(err.NotFound): await message.delete() @@ -179,9 +188,6 @@ class MsG: except exc.MatchError as e: await ctx.send('❌ **No probable match for:** `{}`'.format(e), delete_after=10) - finally: - await asyncio.sleep(self.RATE_LIMIT) - except exc.NotFound: await ctx.send('❌ **No matches found.**', delete_after=10) except ValueError: @@ -202,12 +208,11 @@ class MsG: await ctx.send('✅ **Probable match:**\n{}'.format(await scraper.get_image(post))) + await asyncio.sleep(self.RATE_LIMIT) + except exc.MatchError as e: await ctx.send('❌ **No probable match for:** `{}`'.format(e), delete_after=10) - finally: - await asyncio.sleep(self.RATE_LIMIT) - for attachment in ctx.message.attachments: try: await ctx.trigger_typing() @@ -216,12 +221,11 @@ class MsG: await ctx.send('✅ **Probable match:**\n{}'.format(await scraper.get_image(post))) + await asyncio.sleep(self.RATE_LIMIT) + except exc.MatchError as e: await ctx.send('❌ **No probable match for:** `{}`'.format(e), delete_after=10) - finally: - await asyncio.sleep(self.RATE_LIMIT) - except exc.MissingArgument: await ctx.send('❌ **Invalid url or file.**', delete_after=10) @@ -257,6 +261,8 @@ class MsG: await ctx.send('✅ **Probable match from {}:**\n{}'.format(message.author.display_name, await scraper.get_image(post))) + await asyncio.sleep(self.RATE_LIMIT) + if delete: with suppress(err.NotFound): await message.delete() @@ -264,9 +270,6 @@ class MsG: except exc.MatchError as e: await ctx.send('❌ **No probable match for:** `{}`'.format(e), delete_after=10) - finally: - await asyncio.sleep(self.RATE_LIMIT) - for message in attachments: for attachment in message.attachments: try: @@ -276,6 +279,8 @@ class MsG: await ctx.send('✅ **Probable match from {}:**\n{}'.format(message.author.display_name, await scraper.get_image(post))) + await asyncio.sleep(self.RATE_LIMIT) + if delete: with suppress(err.NotFound): await message.delete() @@ -283,14 +288,88 @@ class MsG: except exc.MatchError as e: await ctx.send('❌ **No probable match for:** `{}`'.format(e), delete_after=10) - finally: - await asyncio.sleep(self.RATE_LIMIT) - except exc.NotFound: await ctx.send('❌ **No matches found.**', delete_after=10) except ValueError: await ctx.send('❌ **Invalid limit.**', delete_after=10) + async def reverse(self): + while self.reversing: + message = await self.queue.get() + + for match in re.finditer('(http[a-z]?:\/\/[^ ]*\.(?:gif|png|jpg|jpeg))', message.content): + try: + await message.channel.trigger_typing() + + post = await scraper.get_post(match.group(0)) + + await message.channel.send('✅ **Probable match from {}:**\n{}'.format(message.author.display_name, await scraper.get_image(post))) + + await asyncio.sleep(self.RATE_LIMIT) + + with suppress(err.NotFound): + await message.delete() + + except exc.MatchError as e: + await message.channel.send('❌ **No probable match for:** `{}`'.format(e), delete_after=10) + + for attachment in message.attachments: + try: + await message.channel.trigger_typing() + + post = await scraper.get_post(attachment.url) + + await message.channel.send('✅ **Probable match from {}:**\n{}'.format(message.author.display_name, await scraper.get_image(post))) + + await asyncio.sleep(self.RATE_LIMIT) + + with suppress(err.NotFound): + await message.delete() + + except exc.MatchError as e: + await message.channel.send('❌ **No probable match for:** `{}`'.format(e), delete_after=10) + + print('STOPPED : reversing') + + async def queue_on_message(self, channel): + def check(msg): + if msg.content.lower() == 'stop' and msg.channel is channel and msg.author.guild_permissions.administrator: + raise exc.Abort + elif msg.channel is channel and msg.author.id != self.bot.user.id and (re.search('(http[a-z]?:\/\/[^ ]*\.(?:gif|png|jpg|jpeg))', msg.content) is not None or msg.attachments): + return True + return False + + try: + while not self.bot.is_closed(): + message = await self.bot.wait_for('message', check=check) + await self.queue.put(message) + + except exc.Abort: + u.tasks['auto_rev'].remove(channel.id) + u.dump(u.tasks, 'cogs/tasks.pkl') + if not u.tasks['auto_rev']: + self.reversing = False + print('STOPPED : looping #{}'.format(channel.name)) + await channel.send('✅ **Stopped queueing messages for reversion in** {}**.**'.format(channel.mention), delete_after=5) + + @commands.command(name='autoreverse', aliases=['autorev', 'ar']) + async def auto_reverse_image_search(self, ctx): + try: + 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_on_message(ctx.channel)) + if not self.reversing: + self.bot.loop.create_task(self.reverse()) + self.reversing = True + print('LOOPING : #{}'.format(ctx.channel.name)) + await ctx.send('✅ **Auto-reversing all images in {}.**'.format(ctx.channel.mention), delete_after=5) + else: + raise exc.Exists + + except exc.Exists: + await ctx.send('❌ **Already auto-reversing in {}.** Type `stop` to stop.'.format(ctx.channel.mention), delete_after=10) + def get_favorites(self, ctx, args): if '-f' in args or '-favs' in args or '-faves' in args or '-favorites' in args: if self.favorites.get(ctx.author.id, {}).get('tags', set()): diff --git a/src/main/cogs/management.py b/src/main/cogs/management.py index 04965d0..3832fe2 100644 --- a/src/main/cogs/management.py +++ b/src/main/cogs/management.py @@ -23,7 +23,7 @@ class Administration: for channel in u.tasks['auto_del']: temp = self.bot.get_channel(channel) self.bot.loop.create_task(self.queue_on_message(temp)) - print('Looping #{}'.format(temp.name)) + print('LOOPING : #{}'.format(temp.name)) self.bot.loop.create_task(self.delete()) self.deleting = True @@ -103,17 +103,18 @@ class Administration: if not message.pinned: await message.delete() + print('STOPPED : deleting') + async def queue_on_message(self, channel): def check(msg): if msg.content.lower() == 'stop' and msg.channel is channel and msg.author.guild_permissions.administrator: raise exc.Abort elif msg.channel is channel and not msg.pinned: return True - else: - return False + return False try: - async for message in channel.history(): + async for message in channel.history(limit=None): if message.content.lower() == 'stop' and message.author.guild_permissions.administrator: raise exc.Abort if not message.pinned: @@ -128,13 +129,9 @@ class Administration: u.dump(u.tasks, 'cogs/tasks.pkl') if not u.tasks['auto_del']: self.deleting = False - print('Stopped deleting.') - print('Stopped looping #{}'.format(channel.name)) + print('STOPPED : looping #{}'.format(channel.name)) await channel.send('✅ **Stopped queueing messages for deletion in** {}**.**'.format(channel.mention), delete_after=5) - except AttributeError: - pass - @commands.command(name='autodelete', aliases=['autodel', 'ad']) @commands.has_permissions(administrator=True) @checks.del_ctx() @@ -147,13 +144,13 @@ class Administration: if not self.deleting: self.bot.loop.create_task(self.delete()) self.deleting = True - print('Looping #{}'.format(ctx.channel.name)) - await ctx.send('✅ **Auto-deleting all messages in this channel.**', delete_after=5) + print('LOOPING : #{}'.format(ctx.channel.name)) + await ctx.send('✅ **Auto-deleting all messages in {}.**'.format(ctx.channel.mention), delete_after=5) else: raise exc.Exists except exc.Exists: - await ctx.send('❌ **Already auto-deleting in this channel.** Type `stop` to stop.', delete_after=10) + await ctx.send('❌ **Already auto-deleting in {}.** Type `stop` to stop.'.format(ctx.channel.mention), delete_after=10) @commands.command(name='deletecommands', aliases=['delcmds']) @commands.has_permissions(administrator=True) diff --git a/src/main/utils/utils.py b/src/main/utils/utils.py index c879555..cba8d02 100644 --- a/src/main/utils/utils.py +++ b/src/main/utils/utils.py @@ -53,7 +53,7 @@ def dump(obj, filename): settings = setdefault('settings.pkl', {'del_ctx': []}) -tasks = setdefault('cogs/tasks.pkl', {'auto_del': []}) +tasks = setdefault('cogs/tasks.pkl', {'auto_del': [], 'auto_rev': []}) session = aiohttp.ClientSession()