diff --git a/.gitignore b/.gitignore index 974375e..776682a 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,7 @@ *.pyo *.pyc *.DS_Store +*.pkl # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/src/main/cogs/booru.py b/src/main/cogs/booru.py index 3b04df0..33de78e 100644 --- a/src/main/cogs/booru.py +++ b/src/main/cogs/booru.py @@ -1,87 +1,48 @@ -import json - -try: - with open('blacklists.json') as infile: - blacklists = json.load(infile) - print('\"blacklists.json\" loaded.') -except FileNotFoundError: - with open('blacklists.json', 'w+') as iofile: - print('Blacklists file not found: \"blacklists.json\" created and loaded.') - json.dump({'global_blacklist': [], 'guild_blacklist': {}, 'user_blacklist': {}}, iofile, indent=4, sort_keys=True) - iofile.seek(0) - blacklists = json.load(iofile) -try: - with open('aliases.json') as infile: - aliases = json.load(infile) - print('\"aliases.json\" loaded.') -except FileNotFoundError: - with open('aliases.json', 'w+') as iofile: - print('Aliases file not found: \"aliases.json\" created and loaded.') - json.dump({'global_blacklist': {}, 'guild_blacklist': {}, 'user_blacklist': {}}, iofile, indent=4, sort_keys=True) - iofile.seek(0) - aliases = json.load(iofile) - import asyncio -import discord as d -import requests +import json import traceback as tb + import discord as d +from discord import errors as err from discord import reaction from discord.ext import commands from discord.ext.commands import errors as errext -from discord import errors as err -from cogs import tools -from misc import checks -from misc import exceptions as exc -from utils import formatter, scraper -from utils import utils as u -HEADERS = {'user-agent': 'Modumind/0.0.1 (Myned)'} +from cogs import tools +from misc import exceptions as exc +from misc import checks +from utils import utils as u +from utils import formatter, scraper # temp_urls = {} + class MsG: def __init__(self, bot): self.bot = bot + self.LIMIT = 100 - # Creates reaction-based paginator for linked pools - @commands.command(brief='e621/e926 Pool selector', description='e621/e926 | NSFW/SFW\nShow pools in a page format', hidden=True) + self.blacklists = u.setdefault( + './cogs/blacklists.pkl', {'global_blacklist': set(), 'guild_blacklist': {}, 'user_blacklist': {}}) + self.aliases = u.setdefault('./cogs/aliases.pkl', {}) + + # Tag search + @commands.command(aliases=['tag', 't'], brief='e621 Tag search', description='e621 | NSFW\nReturn a link search for given tags') @checks.del_ctx() - async def pool(self, ctx, url): - pool_urls = [] - - def check_right(reaction, user): - return user == ctx.message.author and str(reaction.emoji) == '➡️' - def check_left(reaction, user): - return user == ctx.message.author and str(reaction.emoji) == '⬅️' - - try: - pool = scraper.find_pool(url) - for link in pool: - pool_urls.append(scraper.find_image_url('https://e621.net' + link)) - except exc.PostError: - await ctx.send('❌ ' + ctx.message.author.mention + ' **No pool found.**') - except exc.ImageError: - await ctx.send('❌ ' + ctx.message.author.mention + ' **No image found.**') - except Exception: - await ctx.send('{}\n```{}```'.format(exc.base, tb.format_exc(limit=1))) - tb.print_exc() + async def tags(self, ctx, *tags): + await ctx.send('✅ `{}`\nhttps://e621.net/post?tags={}'.format(formatter.tostring(tags), ','.join(tags))) # Tag aliases @commands.command(aliases=['alias', 'a'], brief='e621 Tag aliases', description='e621 | NSFW\nSearch aliases for given tag') @checks.del_ctx() async def aliases(self, ctx, tag): - global HEADERS aliases = [] - try: - alias_request = requests.get('https://e621.net/tag_alias/index.json?aliased_to=' + tag + '&approved=true', headers=HEADERS).json() - for dic in alias_request: - aliases.append(dic['name']) - await ctx.send('✅ `' + tag + '` **aliases:**\n```' + formatter.tostring(aliases) + '```') - except Exception: - await ctx.send('{}\n```{}```'.format(exc.base, tb.format_exc(limit=1))) - tb.print_exc() + + alias_request = await u.fetch('https://e621.net/tag_alias/index.json', params={'aliased_to': tag, 'approved': 'true'}, json=True) + for dic in alias_request: + aliases.append(dic['name']) + await ctx.send('✅ `' + tag + '` **aliases:**\n```' + formatter.tostring(aliases) + '```') # Reverse image searches a linked image using the public iqdb @commands.command(name='reverse', aliases=['rev', 'ris'], brief='e621 Reverse image search', description='e621 | NSFW\nReverse-search an image with given URL') @@ -89,85 +50,355 @@ class MsG: async def reverse_image_search(self, ctx, url): try: await ctx.trigger_typing() - await ctx.send('✅ ' + ctx.message.author.mention + ' **Probable match:**\n' + scraper.check_match('http://iqdb.harry.lu/?url={}'.format(url))) + await ctx.send('✅ ' + ctx.message.author.mention + ' **Probable match:**\n' + await scraper.check_match(url)) except exc.MatchError: await ctx.send('❌ ' + ctx.message.author.mention + ' **No probable match.**', delete_after=10) - except Exception: - await ctx.send('{}\n```{}```'.format(exc.base, tb.format_exc(limit=1))) - tb.print_exc() - @commands.command(name='e621p', aliases=['e6p', '6p']) + async def return_pool(self, *, ctx, booru='e621', query=[]): + def on_message(msg): + if msg.content.lower() == 'cancel' and msg.author is ctx.message.author and msg.channel is ctx.message.channel: + raise exc.Abort + try: + if int(msg.content) <= len(pools) and int(msg.content) > 0 and msg.author is ctx.message.author and msg.channel is ctx.message.channel: + return True + except ValueError: + pass + else: + return False + + posts = {} + pool_id = None + + pools = [] + pool_request = await u.fetch('https://{}.net/pool/index.json'.format(booru), params={'query': ' '.join(query)}, json=True) + if len(pool_request) > 1: + for pool in pool_request: + pools.append(pool['name']) + match = await ctx.send('✅ **Multiple pools found.** Type in the correct match.\n```\n{}```\nor `cancel` to cancel.'.format('\n'.join(['{} {}'.format(c, elem) for c, elem in enumerate(pools, 1)]))) + try: + selection = await self.bot.wait_for('message', check=on_message, timeout=10 * 60) + except exc.Abort: + raise exc.Abort + finally: + await match.delete() + pool = [pool for pool in pool_request if pool['name'] + == pools[int(selection.content) - 1]][0] + await selection.delete() + pool_id = pool['id'] + elif request: + pool = pool_request[0] + pool_id = pool_request[0]['id'] + else: + raise exc.NotFound + + page = 1 + while len(posts) < pool['post_count']: + posts_request = await u.fetch('https://{}.net/pool/show.json'.format(booru), params={'id': pool['id'], 'page': page}, json=True) + for post in posts_request['posts']: + posts[post['id']] = post['file_url'] + page += 1 + + return posts, pool_id + + # Creates reaction-based paginator for linked pools + @commands.command(name='pool', aliases=['e6pp'], brief='e621 pool paginator', description='e621 | NSFW\nShow pools in a page format', hidden=True) @checks.del_ctx() - @checks.is_nsfw() - async def e621_paginator(self, ctx, *args): - def react(reaction, user): - if reaction.emoji == '⬅' and reaction.message.content == paginator.content and user is ctx.message.author: raise exc.Left - elif reaction.emoji == '🚫' and reaction.message.content == paginator.content and user is ctx.message.author: raise exc.Abort - elif reaction.emoji == '📁' and reaction.message.content == paginator.content and user is ctx.message.author: raise exc.Save - elif reaction.emoji == '➡' and reaction.message.content == paginator.content and user is ctx.message.author: raise exc.Right - else: return False + async def pool_paginator(self, ctx, *kwords): + def on_react(reaction, user): + if reaction.emoji == '🚫' and reaction.message.content == paginator.content and user is ctx.message.author: + raise exc.Abort + elif reaction.emoji == '📁' and reaction.message.content == paginator.content and user is ctx.message.author: + raise exc.Save + elif reaction.emoji == '⬅' and reaction.message.content == paginator.content and user is ctx.message.author: + raise exc.Left + elif reaction.emoji == '🔢' and reaction.message.content == paginator.content and user is ctx.message.author: + raise exc.GoTo + elif reaction.emoji == '➡' and reaction.message.content == paginator.content and user is ctx.message.author: + raise exc.Right + else: + return False + + def on_message(msg): + try: + if int(msg.content) <= len(posts) and msg.author is ctx.message.author and msg.channel is ctx.message.channel: + return True + except ValueError: + pass + else: + return False user = ctx.message.author - args = list(args) - limit = 100 + starred = [] + c = 1 try: await ctx.trigger_typing() - c = 1 - posts = self.check_return_posts(ctx=ctx, booru='e621', tags=args, limit=limit) - starred = [] + posts, pool_id = await self.return_pool(ctx=ctx, booru='e621', query=kwords) + keys = list(posts.keys()) + values = list(posts.values()) + + embed = d.Embed( + title='/post/{}'.format(keys[c - 1]), url='https://e621.net/post/show/{}'.format(keys[c - 1]), color=ctx.me.color).set_image(url=values[c - 1]) + embed.set_author(name='/pool/{}'.format(pool_id), + url='https://e621.net/pool/show?id={}'.format(pool_id), icon_url=user.avatar_url) + embed.set_footer(text='{} / {}'.format(c, len(posts)), + icon_url='http://ndl.mgccw.com/mu3/app/20141013/18/1413204353554/icon/icon_xl.png') - embed = d.Embed(title='/post/{}'.format(list(posts.keys())[c-1]), url='https://e621.net/post/show/{}'.format(list(posts.keys())[c-1]), color=ctx.me.color).set_image(url=list(posts.values())[c-1]) - embed.set_author(name=formatter.tostring(args, random=True), url='https://e621.net/post?tags={}'.format(','.join(args)), icon_url=user.avatar_url) - embed.set_footer(text='e621', icon_url='http://ndl.mgccw.com/mu3/app/20141013/18/1413204353554/icon/icon_xl.png') paginator = await ctx.send(embed=embed) - await paginator.add_reaction('⬅') await paginator.add_reaction('🚫') await paginator.add_reaction('📁') + await paginator.add_reaction('⬅') + await paginator.add_reaction('🔢') await paginator.add_reaction('➡') await asyncio.sleep(1) while True: try: - await self.bot.wait_for('reaction_add', check=react, timeout=5*60) + await self.bot.wait_for('reaction_add', check=on_react, timeout=10 * 60) + except exc.Left: if c > 1: c -= 1 - embed.title = '/post/{}'.format(list(posts.keys())[c-1]) - embed.url = 'https://e621.net/post/show/{}'.format(list(posts.keys())[c-1]) - embed.set_image(url=list(posts.values())[c-1]) - await paginator.edit(embed=embed) + embed.title = '/post/{}'.format(keys[c - 1]) + embed.url = 'https://e621.net/post/show/{}'.format(keys[c - 1]) + embed.set_footer(text='{} / {}'.format(c, len(posts)), + icon_url='http://ndl.mgccw.com/mu3/app/20141013/18/1413204353554/icon/icon_xl.png') + embed.set_image(url=values[c - 1]) + + await paginator.edit(content=None, embed=embed) + else: + await paginator.edit(content='❌ **First image.**') + + except exc.GoTo: + await paginator.edit(content='**Enter image number...**') + number = await self.bot.wait_for('message', check=on_message, timeout=10 * 60) + + c = int(number.content) + await number.delete() + embed.title = '/post/{}'.format(keys[c - 1]) + embed.url = 'https://e621.net/post/show/{}'.format(keys[c - 1]) + embed.set_footer(text='{} / {}'.format(c, len(posts)), + icon_url='http://ndl.mgccw.com/mu3/app/20141013/18/1413204353554/icon/icon_xl.png') + embed.set_image(url=values[c - 1]) + + await paginator.edit(content=None, embed=embed) + except exc.Save: - if list(posts.values())[c-1] not in starred: - starred.append(list(posts.values())[c-1]) + if values[c - 1] not in starred: + starred.append(values[c - 1]) + + await paginator.edit(content='**Image** `{}` **saved.**'.format(len(starred))) + + except exc.Right: + if c < len(keys): + c += 1 + embed.title = '/post/{}'.format(keys[c - 1]) + embed.url = 'https://e621.net/post/show/{}'.format(keys[c - 1]) + embed.set_footer(text='{} / {}'.format(c, len(posts)), + icon_url='http://ndl.mgccw.com/mu3/app/20141013/18/1413204353554/icon/icon_xl.png') + embed.set_image(url=values[c - 1]) + + await paginator.edit(content=None, embed=embed) + + except exc.Abort: + try: + await paginator.edit(content='🚫 **Exited paginator.**') + except UnboundLocalError: + await ctx.send('🚫 **Exited paginator.**') + except asyncio.TimeoutError: + try: + await ctx.send(content='❌ **Paginator timed out.**') + except UnboundLocalError: + await ctx.send('❌ **Paginator timed out.**') + except exc.NotFound: + await ctx.send('❌ **Pool not found.**', delete_after=10) + except exc.Timeout: + await ctx.send('❌ **Request timed out.**') + finally: + for url in starred: + await user.send(url) + if len(starred) > 5: + await asyncio.sleep(2.1) + + # Messy code that checks image limit and tags in blacklists + async def check_return_posts(self, *, ctx, booru='e621', tags=[], limit=1, previous=[]): + guild = ctx.message.guild if isinstance( + ctx.message.guild, d.Guild) else ctx.message.channel + channel = ctx.message.channel + user = ctx.message.author + + blacklist = set() + # Creates temp blacklist based on context + for tag in self.blacklists['global_blacklist']: + blacklist.update(list(self.aliases[tag]) + [tag]) + for tag in self.blacklists['guild_blacklist'].get(guild.id, {}).get(channel.id, set()): + blacklist.update(list(self.aliases[tag]) + [tag]) + for tag in self.blacklists['user_blacklist'].get(user.id, set()): + blacklist.update(list(self.aliases[tag]) + [tag]) + # Checks if tags are in local blacklists + if tags: + if len(tags) > 5: + raise exc.TagBoundsError(formatter.tostring(tags[5:])) + for tag in tags: + if tag == 'swf' or tag == 'webm' or tag in blacklist: + raise exc.TagBlacklisted(tag) + + # Checks for blacklisted tags in endpoint blacklists - try/except is for continuing the parent loop + posts = {} + c = 0 + while len(posts) < limit: + if c == 50 + limit * 3: + raise exc.Timeout + request = await u.fetch('https://{}.net/post/index.json'.format(booru), params={'tags': ','.join(['order:random'] + tags), 'limit': self.LIMIT}, json=True) + if len(request) == 0: + raise exc.NotFound(formatter.tostring(tags)) + if len(request) < limit: + limit = len(request) + for post in request: + if 'swf' in post['file_ext'] or 'webm' in post['file_ext']: + continue + try: + for tag in blacklist: + if tag in post['tags']: + raise exc.Continue + except exc.Continue: + continue + if post['file_url'] not in posts.values() and post['file_url'] not in previous: + posts[post['id']] = post['file_url'] + if len(posts) == limit: + break + c += 1 + return posts + + @commands.command(name='e621p', aliases=['e6p', '6p']) + @checks.del_ctx() + @checks.is_nsfw() + async def e621_paginator(self, ctx, *args): + def on_react(reaction, user): + if reaction.emoji == '🚫' and reaction.message.content == paginator.content and user is ctx.message.author: + raise exc.Abort + elif reaction.emoji == '📁' and reaction.message.content == paginator.content and user is ctx.message.author: + raise exc.Save + elif reaction.emoji == '⬅' and reaction.message.content == paginator.content and user is ctx.message.author: + raise exc.Left + elif reaction.emoji == '🔢' and reaction.message.content == paginator.content and user is ctx.message.author: + raise exc.GoTo + elif reaction.emoji == '➡' and reaction.message.content == paginator.content and user is ctx.message.author: + raise exc.Right + else: + return False + + def on_message(msg): + try: + if int(msg.content) <= len(posts) and msg.author is ctx.message.author and msg.channel is ctx.message.channel: + return True + except ValueError: + pass + else: + return False + + user = ctx.message.author + args = list(args) + limit = self.LIMIT / 5 + starred = [] + c = 1 + + try: + await ctx.trigger_typing() + + posts = await self.check_return_posts(ctx=ctx, booru='e621', tags=args, limit=limit) + keys = list(posts.keys()) + values = list(posts.values()) + + embed = d.Embed( + title='/post/{}'.format(keys[c - 1]), url='https://e621.net/post/show/{}'.format(keys[c - 1]), color=ctx.me.color).set_image(url=values[c - 1]) + embed.set_author(name=formatter.tostring(args, random=True), + url='https://e621.net/post?tags={}'.format(','.join(args)), icon_url=user.avatar_url) + embed.set_footer(text='{} / {}'.format(c, len(posts)), + icon_url='http://ndl.mgccw.com/mu3/app/20141013/18/1413204353554/icon/icon_xl.png') + + paginator = await ctx.send(embed=embed) + + await paginator.add_reaction('🚫') + await paginator.add_reaction('📁') + await paginator.add_reaction('⬅') + await paginator.add_reaction('🔢') + await paginator.add_reaction('➡') + await asyncio.sleep(1) + + while True: + try: + await self.bot.wait_for('reaction_add', check=on_react, timeout=10 * 60) + + except exc.Left: + if c > 1: + c -= 1 + embed.title = '/post/{}'.format(keys[c - 1]) + embed.url = 'https://e621.net/post/show/{}'.format(keys[c - 1]) + embed.set_footer(text='{} / {}'.format(c, len(posts)), + icon_url='http://ndl.mgccw.com/mu3/app/20141013/18/1413204353554/icon/icon_xl.png') + embed.set_image(url=values[c - 1]) + await paginator.edit(content=None, embed=embed) + else: + await paginator.edit(content='❌ **First image.**') + + except exc.GoTo: + await paginator.edit(content='**Enter image number...**') + number = await self.bot.wait_for('message', check=on_message, timeout=10 * 60) + + c = int(number.content) + await number.delete() + embed.title = '/post/{}'.format(keys[c - 1]) + embed.url = 'https://e621.net/post/show/{}'.format(keys[c - 1]) + embed.set_footer(text='{} / {}'.format(c, len(posts)), + icon_url='http://ndl.mgccw.com/mu3/app/20141013/18/1413204353554/icon/icon_xl.png') + embed.set_image(url=values[c - 1]) + + await paginator.edit(content=None, embed=embed) + + except exc.Save: + if values[c - 1] not in starred: + starred.append(values[c - 1]) + + await paginator.edit(content='**Image** `{}` **saved.**'.format(len(starred))) + except exc.Right: if c % limit == 0: await ctx.trigger_typing() - try: posts.update(self.check_return_posts(ctx=ctx, booru='e621', tags=args, limit=limit, previous=posts)) + try: + posts.update(await self.check_return_posts(ctx=ctx, booru='e621', tags=args, limit=limit, previous=posts)) except exc.NotFound: await paginator.edit(content='❌ **No more images found.**') + keys = list(posts.keys()) + values = list(posts.values()) c += 1 - embed.title = '/post/{}'.format(list(posts.keys())[c-1]) - embed.url = 'https://e621.net/post/show/{}'.format(list(posts.keys())[c-1]) - embed.set_image(url=list(posts.values())[c-1]) - await paginator.edit(embed=embed) - except exc.Abort: await paginator.edit(content='🚫 **Exited paginator.**') - except exc.TagBlacklisted as e: await ctx.send('❌ `{}` **blacklisted.**'.format(e), delete_after=10) - except exc.TagBoundsError as e: await ctx.send('❌ `{}` **out of bounds.** Tags limited to 5, currently.'.format(e), delete_after=10) - except exc.Timeout: await ctx.send('❌ **Request timed out.**') - except asyncio.TimeoutError: await paginator.edit(content='❌ **Paginator timed out.**') - except Exception: - await ctx.send('{}\n```{}```'.format(exc.base, tb.format_exc(limit=1))) - tb.print_exc() + embed.title = '/post/{}'.format(keys[c - 1]) + embed.url = 'https://e621.net/post/show/{}'.format(keys[c - 1]) + embed.set_footer(text='{} / {}'.format(c, len(posts)), + icon_url='http://ndl.mgccw.com/mu3/app/20141013/18/1413204353554/icon/icon_xl.png') + embed.set_image(url=values[c - 1]) + await paginator.edit(content=None, embed=embed) + + except exc.Abort: + await paginator.edit(content='🚫 **Exited paginator.**') + except asyncio.TimeoutError: + await paginator.edit(content='❌ **Paginator timed out.**') + except exc.NotFound as e: + await ctx.send('❌ `{}` **not found.**'.format(e), delete_after=10) + except exc.TagBlacklisted as e: + await ctx.send('❌ `{}` **blacklisted.**'.format(e), delete_after=10) + except exc.TagBoundsError as e: + await ctx.send('❌ `{}` **out of bounds.** Tags limited to 5, currently.'.format(e), delete_after=10) + except exc.Timeout: + await ctx.send('❌ **Request timed out.**') finally: - if starred: - for url in starred: - await user.send(url) - if len(starred) > 5: - await asyncio.sleep(2.1) + for url in starred: + await user.send(url) + if len(starred) > 5: + await asyncio.sleep(2.1) @e621_paginator.error async def e621_paginator_error(self, ctx, error): @@ -192,23 +423,31 @@ class MsG: if int(arg) <= 6 and int(arg) >= 1: limit = int(arg) args.remove(arg) - else: raise exc.BoundsError(arg) - posts = self.check_return_posts(ctx=ctx, booru='e621', tags=args, limit=limit)#, previous=temp_urls.get(ctx.message.author.id, [])) + else: + raise exc.BoundsError(arg) + # , previous=temp_urls.get(ctx.message.author.id, [])) + posts = await self.check_return_posts(ctx=ctx, booru='e621', tags=args, limit=limit) for ident, url in posts.items(): - embed = d.Embed(title='/post/{}'.format(ident), url='https://e621.net/post/show/{}'.format(ident), color=ctx.me.color).set_image(url=url) - embed.set_author(name=formatter.tostring(args, random=True), url='https://e621.net/post?tags={}'.format(','.join(args)), icon_url=ctx.message.author.avatar_url) - embed.set_footer(text='e621', icon_url='http://ndl.mgccw.com/mu3/app/20141013/18/1413204353554/icon/icon_xl.png') + embed = d.Embed(title='/post/{}'.format(ident), url='https://e621.net/post/show/{}'.format(ident), + color=ctx.me.color).set_image(url=url) + embed.set_author(name=formatter.tostring(args, random=True), + url='https://e621.net/post?tags={}'.format(','.join(args)), icon_url=ctx.message.author.avatar_url) + embed.set_footer( + text='e621', icon_url='http://ndl.mgccw.com/mu3/app/20141013/18/1413204353554/icon/icon_xl.png') await ctx.send(embed=embed) # temp_urls.setdefault(ctx.message.author.id, []).extend(posts.values()) - except exc.TagBlacklisted as e: await ctx.send('❌ `' + str(e) + '` **blacklisted.**', delete_after=10) - except exc.BoundsError as e: await ctx.send('❌ `' + str(e) + '` **out of bounds.**', delete_after=10) - except exc.TagBoundsError as e: await ctx.send('❌ `' + str(e) + '` **out of bounds.** Tags limited to 5, currently.', delete_after=10) - except exc.NotFound as e: await ctx.send('❌ `' + str(e) + '` **not found.**', delete_after=10) - except exc.Timeout: await ctx.send('❌ **Request timed out.**') - except Exception: - await ctx.send('{}\n```{}```'.format(exc.base, tb.format_exc(limit=1))) - tb.print_exc() - tools.command_dict.setdefault(str(ctx.message.author.id), {}).update({'command': ctx.command, 'args': ctx.args}) + except exc.TagBlacklisted as e: + await ctx.send('❌ `' + str(e) + '` **blacklisted.**', delete_after=10) + except exc.BoundsError as e: + await ctx.send('❌ `' + str(e) + '` **out of bounds.**', delete_after=10) + except exc.TagBoundsError as e: + await ctx.send('❌ `' + str(e) + '` **out of bounds.** Tags limited to 5, currently.', delete_after=10) + except exc.NotFound as e: + await ctx.send('❌ `' + str(e) + '` **not found.**', delete_after=10) + except exc.Timeout: + await ctx.send('❌ **Request timed out.**') + tools.command_dict.setdefault(str(ctx.message.author.id), {}).update( + {'command': ctx.command, 'args': ctx.args}) @e621.error async def e621_error(self, ctx, error): @@ -232,61 +471,29 @@ class MsG: if int(arg) <= 6 and int(arg) >= 1: limit = int(arg) args.remove(arg) - else: raise exc.BoundsError(arg) - posts = self.check_return_posts(ctx=ctx, booru='e926', tags=args, limit=limit)#, previous=temp_urls.get(ctx.message.author.id, [])) + else: + raise exc.BoundsError(arg) + # , previous=temp_urls.get(ctx.message.author.id, [])) + posts = await self.check_return_posts(ctx=ctx, booru='e926', tags=args, limit=limit) for ident, url in posts.items(): - embed = d.Embed(title='/post/{}'.format(ident), url='https://e926.net/post/show/{}'.format(ident), color=ctx.me.color).set_image(url=url) - embed.set_author(name=formatter.tostring(args, random=True), url='https://e621.net/post?tags={}'.format(','.join(args)), icon_url=ctx.message.author.avatar_url) - embed.set_footer(text='e926', icon_url='http://ndl.mgccw.com/mu3/app/20141013/18/1413204353554/icon/icon_xl.png') + embed = d.Embed(title='/post/{}'.format(ident), url='https://e926.net/post/show/{}'.format(ident), + color=ctx.me.color).set_image(url=url) + embed.set_author(name=formatter.tostring(args, random=True), + url='https://e621.net/post?tags={}'.format(','.join(args)), icon_url=ctx.message.author.avatar_url) + embed.set_footer( + text='e926', icon_url='http://ndl.mgccw.com/mu3/app/20141013/18/1413204353554/icon/icon_xl.png') await ctx.send(embed=embed) # temp_urls.setdefault(ctx.message.author.id, []).extend(posts.values()) - except exc.TagBlacklisted as e: await ctx.send('❌ `' + str(e) + '` **blacklisted.**', delete_after=10) - except exc.BoundsError as e: await ctx.send('❌ `' + str(e) + '` **out of bounds.**', delete_after=10) - except exc.TagBoundsError as e: await ctx.send('❌ `' + str(e) + '` **out of bounds.** Tags limited to 5, currently.', delete_after=10) - except exc.NotFound as e: await ctx.send('❌ `' + str(e) + '` **not found.**', delete_after=10) - except exc.Timeout: await ctx.send('❌ **Request timed out.**') - except Exception: - await ctx.send('{}\n```{}```'.format(exc.base, tb.format_exc(limit=1))) - tb.print_exc() - - # Messy code that checks image limit and tags in blacklists - def check_return_posts(self, *, ctx, booru='e621', tags=[], limit=1, previous=[]): - global blacklists, aliases, HEADERS - - if isinstance(ctx.message.guild, d.Guild): guild = ctx.message.guild - else: guild = ctx.message.channel - channel = ctx.message.channel - user = ctx.message.author - blacklist = [] - - # Creates temp blacklist based on context - for k, v in aliases['global_blacklist'].items(): blacklist.extend([k] + v) - for k, v in aliases['guild_blacklist'].get(str(guild.id), {}).get(str(channel.id), {}).items(): blacklist.extend([k] + v) - for k, v in aliases['user_blacklist'].get(str(user.id), {}).items(): blacklist.extend([k] + v) - # Checks if tags are in local blacklists - if tags: - if len(tags) > 5: raise exc.TagBoundsError(formatter.tostring(tags[5:])) - for tag in tags: - if tag == 'swf' or tag == 'webm' or tag in blacklist: raise exc.TagBlacklisted(tag) - - # Checks for blacklisted tags in endpoint blacklists - try/except is for continuing the parent loop - posts = {} - c = 0 - while len(posts) < limit: - if c == 50 + limit: raise exc.Timeout - request = requests.get('https://{}.net/post/index.json?tags={}'.format(booru, ','.join(['order:random'] + tags)), headers=HEADERS).json() - if len(request) == 0: raise exc.NotFound(formatter.tostring(tags)) - if len(request) < limit: limit = len(request) - for post in request: - if 'swf' in post['file_ext'] or 'webm' in post['file_ext']: continue - try: - for tag in blacklist: - if tag in post['tags']: raise exc.Continue - except exc.Continue: continue - if post['file_url'] not in posts.values() and post['file_url'] not in previous: posts[post['id']] = post['file_url'] - if len(posts) == limit: break - c += 1 - return posts + except exc.TagBlacklisted as e: + await ctx.send('❌ `' + str(e) + '` **blacklisted.**', delete_after=10) + except exc.BoundsError as e: + await ctx.send('❌ `' + str(e) + '` **out of bounds.**', delete_after=10) + except exc.TagBoundsError as e: + await ctx.send('❌ `' + str(e) + '` **out of bounds.** Tags limited to 5, currently.', delete_after=10) + except exc.NotFound as e: + await ctx.send('❌ `' + str(e) + '` **not found.**', delete_after=10) + except exc.Timeout: + await ctx.send('❌ **Request timed out.**') # Umbrella command structure to manage global, channel, and user blacklists @commands.group(aliases=['bl', 'b'], brief='Manage blacklists', description='Blacklist base command for managing blacklists\n\n`bl get [blacklist]` to show a blacklist\n`bl set [blacklist] [tags]` to replace a blacklist\n`bl clear [blacklist]` to clear a blacklist\n`bl add [blacklist] [tags]` to add tags to a blacklist\n`bl remove [blacklist] [tags]` to remove tags from a blacklist', usage='[flag] [blacklist] ([tags])') @@ -294,6 +501,7 @@ class MsG: async def blacklist(self, ctx): if ctx.invoked_subcommand is None: await ctx.send('❌ **Use a flag to manage blacklists.**\n*Type* `' + ctx.prefix + 'help bl` *for more info.*', delete_after=10) + @blacklist.error async def blacklist_error(self, ctx, error): if isinstance(error, commands.CheckFailure): @@ -304,244 +512,173 @@ class MsG: @blacklist.group(name='get', aliases=['g']) async def _get_blacklist(self, ctx): if ctx.invoked_subcommand is None: - await ctx.send('❌ **Invalid blacklist.**') + await ctx.send('❌ **Invalid blacklist.**', delete_after=10) + @_get_blacklist.command(name='global', aliases=['gl', 'g']) async def __get_global_blacklist(self, ctx): - global blacklists - await ctx.send('🚫 **Global blacklist:**\n```' + formatter.tostring(blacklists['global_blacklist']) + '```') + await ctx.send('🚫 **Global blacklist:**\n```\n' + formatter.tostring(self.blacklists['global_blacklist']) + '```') + @_get_blacklist.command(name='channel', aliases=['ch', 'c']) async def __get_channel_blacklist(self, ctx): - global blacklists - if isinstance(ctx.message.guild, d.Guild): - guild = ctx.message.guild - else: - guild = ctx.message.channel + guild = ctx.message.guild if isinstance( + ctx.message.guild, d.Guild) else ctx.message.channel channel = ctx.message.channel - await ctx.send('🚫 <#' + str(channel.id) + '> **blacklist:**\n```' + formatter.tostring(blacklists['guild_blacklist'].get(str(guild.id), {}).get(str(channel.id), [])) + '```') + await ctx.send('🚫 <#' + channel.id + '> **blacklist:**\n```\n' + formatter.tostring(self.blacklists['guild_blacklist'].get(guild.id, {}).get(channel.id, set())) + '```') + @_get_blacklist.command(name='me', aliases=['m']) async def __get_user_blacklist(self, ctx): - global blacklists user = ctx.message.author - await ctx.send('🚫 ' + user.mention + '**\'s blacklist:**\n```' + formatter.tostring(blacklists['user_blacklist'].get(str(user.id), [])) + '```', delete_after=10) + await ctx.send('🚫 ' + user.mention + '**\'s blacklist:**\n```\n' + formatter.tostring(self.blacklists['user_blacklist'].get(user.id, set())) + '```', delete_after=10) + @_get_blacklist.command(name='here', aliases=['h']) async def __get_here_blacklists(self, ctx): - global blacklists - if isinstance(ctx.message.guild, d.Guild): - guild = ctx.message.guild - else: - guild = ctx.message.channel + guild = ctx.message.guild if isinstance( + ctx.message.guild, d.Guild) else ctx.message.channel channel = ctx.message.channel - await ctx.send('🚫 **__Blacklisted:__**\n\n**Global:**\n```' + formatter.tostring(blacklists['global_blacklist']) + '```\n**<#' + str(channel.id) + '>:**\n```' + formatter.tostring(blacklists['guild_blacklist'].get(str(guild.id), {}).get(str(channel.id), [])) + '```') + await ctx.send('🚫 **__Blacklisted:__**\n\n**Global:**\n```\n' + formatter.tostring(self.blacklists['global_blacklist']) + '```\n**<#' + channel.id + '>:**\n```\n' + formatter.tostring(self.blacklists['guild_blacklist'].get(guild.id, {}).get(channel.id, set())) + '```') + @_get_blacklist.group(name='all', aliases=['a']) async def __get_all_blacklists(self, ctx): if ctx.invoked_subcommand is None: await ctx.send('❌ **Invalid blacklist.**') + @__get_all_blacklists.command(name='guild', aliases=['g']) @commands.has_permissions(manage_channels=True) async def ___get_all_guild_blacklists(self, ctx): - global blacklists - if isinstance(ctx.message.guild, d.Guild): - guild = ctx.message.guild - else: - guild = ctx.message.channel - await ctx.send('🚫 **__' + guild.name + ' blacklists:__**\n\n' + formatter.dict_tostring(blacklists['guild_blacklist'].get(str(guild.id), {}))) + guild = ctx.message.guild if isinstance( + ctx.message.guild, d.Guild) else ctx.message.channel + await ctx.send('🚫 **__' + guild.name + ' blacklists:__**\n\n' + formatter.dict_tostring(self.blacklists['guild_blacklist'].get(guild.id, {}))) + @__get_all_blacklists.command(name='user', aliases=['u', 'member', 'm']) @commands.is_owner() async def ___get_all_user_blacklists(self, ctx): - global blacklists - await ctx.send('🚫 **__User blacklists:__**\n\n' + formatter.dict_tostring(blacklists['user_blacklist'])) + await ctx.send('🚫 **__User blacklists:__**\n\n' + formatter.dict_tostring(self.blacklists['user_blacklist'])) @blacklist.group(name='add', aliases=['a']) async def _add_tags(self, ctx): if ctx.invoked_subcommand is None: - await ctx.send('❌ **Invalid blacklist.**') + await ctx.send('❌ **Invalid blacklist.**', delete_after=10) + @_add_tags.command(name='global', aliases=['gl', 'g']) @commands.is_owner() async def __add_global_tags(self, ctx, *tags): - global blacklists, aliases, HEADERS - try: - for tag in tags: - if tag in blacklists['global_blacklist']: - raise exc.TagExists(tag) - blacklists['global_blacklist'].extend(tags) - for tag in tags: - alias_request = requests.get('https://e621.net/tag_alias/index.json?aliased_to=' + tag + '&approved=true', headers=HEADERS).json() + self.blacklists['global_blacklist'].update(tags) + for tag in tags: + alias_request = await u.fetch('https://e621.net/tag_alias/index.json', params={'aliased_to': tag, 'approved': 'true'}, json=True) + if alias_request: for dic in alias_request: - aliases['global_blacklist'].setdefault(tag, []).append(dic['name']) - with open('blacklists.json', 'w') as outfile: - json.dump(blacklists, outfile, indent=4, sort_keys=True) - with open('aliases.json', 'w') as outfile: - json.dump(aliases, outfile, indent=4, sort_keys=True) - await ctx.send('✅ **Added to global blacklist:**\n```' + formatter.tostring(tags) + '```', delete_after=5) - except exc.TagExists as e: - await ctx.send('❌ `' + str(e) + '` **already in blacklist.**', delete_after=10) - except Exception: - await ctx.send('{}\n```{}```'.format(exc.base, tb.format_exc(limit=1))) - tb.print_exc() + self.aliases.setdefault(tag, set()).add(dic['name']) + u.dump(self.blacklists, './cogs/blacklists.pkl') + u.dump(self.aliases, './cogs/aliases.pkl') + await ctx.send('✅ **Added to global blacklist:**\n```\n' + formatter.tostring(tags) + '```', delete_after=5) + @_add_tags.command(name='channel', aliases=['ch', 'c']) @commands.has_permissions(manage_channels=True) async def __add_channel_tags(self, ctx, *tags): - global blacklists, aliases, HEADERS - if isinstance(ctx.message.guild, d.Guild): - guild = ctx.message.guild - else: - guild = ctx.message.channel + guild = ctx.message.guild if isinstance( + ctx.message.guild, d.Guild) else ctx.message.channel channel = ctx.message.channel - try: - for tag in tags: - if tag in blacklists['guild_blacklist'].get(str(guild.id), {}).get(str(channel.id), []): - raise exc.TagExists(tag) - blacklists['guild_blacklist'].setdefault(str(guild.id), {}).setdefault(str(channel.id), []).extend(tags) - for tag in tags: - alias_request = requests.get('https://e621.net/tag_alias/index.json?aliased_to=' + tag + '&approved=true', headers=HEADERS).json() + + self.blacklists['guild_blacklist'].setdefault( + guild.id, {}).setdefault(channel.id, set()).update(tags) + for tag in tags: + alias_request = await u.fetch('https://e621.net/tag_alias/index.json', params={'aliased_to': tag, 'approved': 'true'}, json=True) + if alias_request: for dic in alias_request: - aliases['guild_blacklist'].setdefault(str(guild.id), {}).setdefault(str(channel.id), {}).setdefault(tag, []).append(dic['name']) - with open('blacklists.json', 'w') as outfile: - json.dump(blacklists, outfile, indent=4, sort_keys=True) - with open('aliases.json', 'w') as outfile: - json.dump(aliases, outfile, indent=4, sort_keys=True) - await ctx.send('✅ **Added to** <#' + str(channel.id) + '> **blacklist:**\n```' + formatter.tostring(tags) + '```', delete_after=5) - except exc.TagExists as e: - await ctx.send('❌ `' + str(e) + '` **already in blacklist.**', delete_after=10) - except Exception: - await ctx.send('{}\n```{}```'.format(exc.base, tb.format_exc(limit=1))) - tb.print_exc() + self.aliases.setdefault(tag, set()).add(dic['name']) + u.dump(self.blacklists, './cogs/blacklists.pkl') + u.dump(self.aliases, './cogs/aliases.pkl') + await ctx.send('✅ **Added to** <#' + channel.id + '> **blacklist:**\n```\n' + formatter.tostring(tags) + '```', delete_after=5) + @_add_tags.command(name='me', aliases=['m']) async def __add_user_tags(self, ctx, *tags): - global blacklists, aliases, HEADERS user = ctx.message.author - try: - for tag in tags: - if tag in blacklists['user_blacklist'].get(str(user.id), []): - raise exc.TagExists(tag) - blacklists['user_blacklist'].setdefault(str(user.id), []).extend(tags) - for tag in tags: - alias_request = requests.get('https://e621.net/tag_alias/index.json?aliased_to=' + tag + '&approved=true', headers=HEADERS).json() + + self.blacklists['user_blacklist'].setdefault(user.id, set()).update(tags) + for tag in tags: + alias_request = await u.fetch('https://e621.net/tag_alias/index.json', params={'aliased_to': tag, 'approved': 'true'}, json=True) + if alias_request: for dic in alias_request: - aliases['user_blacklist'].setdefault(str(user.id), {}).setdefault(tag, []).append(dic['name']) - with open('blacklists.json', 'w') as outfile: - json.dump(blacklists, outfile, indent=4, sort_keys=True) - with open('aliases.json', 'w') as outfile: - json.dump(aliases, outfile, indent=4, sort_keys=True) - await ctx.send('✅ ' + user.mention + ' **added:**\n```' + formatter.tostring(tags) + '```', delete_after=5) - except exc.TagExists as e: - await ctx.send('❌ `' + str(e) + '` **already in blacklist.**', delete_after=10) - except Exception: - await ctx.send('{}\n```{}```'.format(exc.base, tb.format_exc(limit=1))) - tb.print_exc() + self.aliases.setdefault(tag, set()).add(dic['name']) + u.dump(self.blacklists, './cogs/blacklists.pkl') + u.dump(self.aliases, './cogs/aliases.pkl') + await ctx.send('✅ ' + user.mention + ' **added:**\n```\n' + formatter.tostring(tags) + '```', delete_after=5) @blacklist.group(name='remove', aliases=['rm', 'r']) async def _remove_tags(self, ctx): if ctx.invoked_subcommand is None: - await ctx.send('❌ **Invalid blacklist.**') + await ctx.send('❌ **Invalid blacklist.**', delete_after=10) + @_remove_tags.command(name='global', aliases=['gl', 'g']) @commands.is_owner() async def __remove_global_tags(self, ctx, *tags): - global blacklists, aliases try: for tag in tags: - if tag in blacklists['global_blacklist']: - blacklists['global_blacklist'].remove(tag) - del aliases['global_blacklist'][tag] - else: + try: + self.blacklists['global_blacklist'].remove(tag) + except KeyError: raise exc.TagError(tag) - with open('blacklists.json', 'w') as outfile: - json.dump(blacklists, outfile, indent=4, sort_keys=True) - with open('aliases.json', 'w') as outfile: - json.dump(aliases, outfile, indent=4, sort_keys=True) - await ctx.send('✅ **Removed from global blacklist:**\n```' + formatter.tostring(tags) + '```', delete_after=5) + u.dump(self.blacklists, './cogs/blacklists.pkl') + await ctx.send('✅ **Removed from global blacklist:**\n```\n' + formatter.tostring(tags) + '```', delete_after=5) except exc.TagError as e: await ctx.send('❌ `' + str(e) + '` **not in blacklist.**', delete_after=10) - except Exception: - await ctx.send('{}\n```{}```'.format(exc.base, tb.format_exc(limit=1))) - tb.print_exc() + @_remove_tags.command(name='channel', aliases=['ch', 'c']) @commands.has_permissions(manage_channels=True) async def __remove_channel_tags(self, ctx, *tags): - global blacklists, aliases - if isinstance(ctx.message.guild, d.Guild): - guild = ctx.message.guild - else: - guild = ctx.message.channel + guild = ctx.message.guild if isinstance( + ctx.message.guild, d.Guild) else ctx.message.channel channel = ctx.message.channel try: for tag in tags: - if tag in blacklists['guild_blacklist'][str(guild.id)][str(channel.id)]: - blacklists['guild_blacklist'][str(guild.id)][str(channel.id)].remove(tag) - del aliases['guild_blacklist'][str(guild.id)][str(channel.id)][tag] - else: + try: + self.blacklists['guild_blacklist'][guild.id][channel.id].remove(tag) + except KeyError: raise exc.TagError(tag) - with open('blacklists.json', 'w') as outfile: - json.dump(blacklists, outfile, indent=4, sort_keys=True) - with open('aliases.json', 'w') as outfile: - json.dump(aliases, outfile, indent=4, sort_keys=True) - await ctx.send('✅ **Removed from** <#' + str(channel.id) + '> **blacklist:**\n```' + formatter.tostring(tags) + '```', delete_after=5) + u.dump(self.blacklists, './cogs/blacklists.pkl') + await ctx.send('✅ **Removed from** <#' + channel.id + '> **blacklist:**\n```\n' + formatter.tostring(tags) + '```', delete_after=5) except exc.TagError as e: await ctx.send('❌ `' + str(e) + '` **not in blacklist.**', delete_after=10) - except Exception: - await ctx.send('{}\n```{}```'.format(exc.base, tb.format_exc(limit=1))) - tb.print_exc() + @_remove_tags.command(name='me', aliases=['m']) async def __remove_user_tags(self, ctx, *tags): - global blacklists, aliases user = ctx.message.author try: for tag in tags: - if tag in blacklists['user_blacklist'][str(user.id)]: - blacklists['user_blacklist'][str(user.id)].remove(tag) - del aliases['user_blacklist'][str(user.id)][tag] - else: + try: + self.blacklists['user_blacklist'][user.id].remove(tag) + except KeyError: raise exc.TagError(tag) - with open('blacklists.json', 'w') as outfile: - json.dump(blacklists, outfile, indent=4, sort_keys=True) - with open('aliases.json', 'w') as outfile: - json.dump(aliases, outfile, indent=4, sort_keys=True) - await ctx.send('✅ ' + user.mention + ' **removed:**\n```' + formatter.tostring(tags) + '```', delete_after=5) + u.dump(self.blacklists, './cogs/blacklists.pkl') + await ctx.send('✅ ' + user.mention + ' **removed:**\n```\n' + formatter.tostring(tags) + '```', delete_after=5) except exc.TagError as e: await ctx.send('❌ `' + str(e) + '` **not in blacklist.**', delete_after=10) - except Exception: - await ctx.send('{}\n```{}```'.format(exc.base, tb.format_exc(limit=1))) - tb.print_exc() @blacklist.group(name='clear', aliases=['cl', 'c']) async def _clear_blacklist(self, ctx): if ctx.invoked_subcommand is None: - await ctx.send('❌ **Invalid blacklist.**') + await ctx.send('❌ **Invalid blacklist.**', delete_after=10) + @_clear_blacklist.command(name='global', aliases=['gl', 'g']) @commands.is_owner() async def __clear_global_blacklist(self, ctx): - global blacklists, aliases - del blacklists['global_blacklist'] - del aliases['global_blacklist'] - with open('blacklists.json', 'w') as outfile: - json.dump(blacklists, outfile, indent=4, sort_keys=True) - with open('aliases.json', 'w') as outfile: - json.dump(aliases, outfile, indent=4, sort_keys=True) + del self.blacklists['global_blacklist'] + u.dump(self.blacklists, './cogs/blacklists.pkl') await ctx.send('✅ **Global blacklist cleared.**', delete_after=5) + @_clear_blacklist.command(name='channel', aliases=['ch', 'c']) @commands.has_permissions(manage_channels=True) async def __clear_channel_blacklist(self, ctx): - global blacklists, aliases - if isinstance(ctx.message.guild, d.Guild): - guild = ctx.message.guild - else: - guild = ctx.message.channel + guild = ctx.message.guild if isinstance( + ctx.message.guild, d.Guild) else ctx.message.channel channel = ctx.message.channel - del blacklists['guild_blacklist'][str(guild.id)][str(channel.id)] - del aliases['guild_blacklist'][str(guild.id)][str(channel.id)] - with open('blacklists.json', 'w') as outfile: - json.dump(blacklists, outfile, indent=4, sort_keys=True) - with open('aliases.json', 'w') as outfile: - json.dump(aliases, outfile, indent=4, sort_keys=True) - await ctx.send('✅ <#' + str(channel.id) + '> **blacklist cleared.**', delete_after=5) + del self.blacklists['guild_blacklist'][str(guild.id)][channel.id] + u.dump(self.blacklists, './cogs/blacklists.pkl') + await ctx.send('✅ <#' + channel.id + '> **blacklist cleared.**', delete_after=5) + @_clear_blacklist.command(name='me', aliases=['m']) async def __clear_user_blacklist(self, ctx): - global blacklists, aliases user = ctx.message.author - del blacklists['user_blacklist'][str(user.id)] - del aliases['user_blacklist'][str(user.id)] - with open('blacklists.json', 'w') as outfile: - json.dump(blacklists, outfile, indent=4, sort_keys=True) - with open('aliases.json', 'w') as outfile: - json.dump(aliases, outfile, indent=4, sort_keys=True) + del self.blacklists['user_blacklist'][user.id] + u.dump(self.blacklists, './cogs/blacklists.pkl') await ctx.send('✅ ' + user.mention + '**\'s blacklist cleared.**', delete_after=5) diff --git a/src/main/cogs/info.py b/src/main/cogs/info.py index d86435b..bea069c 100644 --- a/src/main/cogs/info.py +++ b/src/main/cogs/info.py @@ -1,7 +1,9 @@ import asyncio +import traceback as tb + import discord -import traceback from discord.ext import commands + from misc import exceptions as exc @@ -10,6 +12,19 @@ class Info: def __init__(self, bot): self.bot = bot + @commands.command(hidden=True) + async def hi(ctx): + user = ctx.message.author + + hello = 'Hewwo, {}.'.format(user.mention) + if user.id == checks.owner_id: + hello += '.. ***Master.*** uwu' + elif user.guild_permissions.administrator: + hello = '{} **Admin** {}'.format(hello[:7], hello[7:]) + elif user.guild_permissions.ban_members: + hello = '{} **Mod** {}'.format(hello[:7], hello[7:]) + await ctx.send(hello) + @commands.group(name='info', aliases=['i']) async def info(self, ctx): if invoked_subcommand is None: @@ -17,16 +32,8 @@ class Info: @info.command(aliases=['g', 'server', 's'], brief='Provides info about a guild', hidden=True) async def guild(self, ctx): - try: - guild = '' - except Exception: - await ctx.send(exc.base) - traceback.print_exc(limit=1) + pass @info.command(aliases=['u', 'member', 'm'], brief='Provides info about a user', hidden=True) async def user(self, ctx): - try: - user = '' - except Exception: - await ctx.send(exc.base) - traceback.print_exc(limit=1) + pass diff --git a/src/main/cogs/management.py b/src/main/cogs/management.py index 36aed10..642d747 100644 --- a/src/main/cogs/management.py +++ b/src/main/cogs/management.py @@ -1,13 +1,16 @@ import asyncio +import traceback as tb + import discord as d -import traceback from discord.ext import commands -from misc import checks + from misc import exceptions as exc +from misc import checks from utils import utils as u RATE_LIMIT = 2.1 + class Administration: def __init__(self, bot): @@ -94,7 +97,7 @@ class Administration: history.extend(await channel.history(limit=None).flatten()) await ch_sent.edit(content='🗄 **Cached** `{}/{}` **channels.**'.format(channels.index(channel) + 1, len(channels))) await asyncio.sleep(RATE_LIMIT) - elif when =='before': + elif when == 'before': for channel in channels: history.extend(await channel.history(limit=None, before=ref.created_at).flatten()) await ch_sent.edit(content='🗄 **Cached** `{}/{}` **channels.**'.format(channels.index(channel) + 1, len(channels))) @@ -117,8 +120,10 @@ class Administration: await cont_sent.delete() del_sent = await ctx.send('🗑 **Deleting messages...**') for message in history: - try: await message.delete() - except d.NotFound: pass + try: + await message.delete() + except d.NotFound: + pass # print('Deleted {}/{} messages.'.format(history.index(message) + 1, len(history))) await del_sent.edit(content='🗑 **Deleted** `{}/{}` **messages.**'.format(history.index(message) + 1, len(history))) await asyncio.sleep(RATE_LIMIT) @@ -127,20 +132,19 @@ class Administration: await ctx.send('❌ **Deletion aborted.**', delete_after=10) except TimeoutError: await ctx.send('❌ **Deletion timed out.**', delete_after=10) - except Exception: - await ctx.send('{}\n```{}```'.format(exc.base, traceback.format_exc(limit=1))) - traceback.print_exc() async def delete(self): while True: message = await self.queue.get() await asyncio.sleep(RATE_LIMIT) - try: await message.delete() - except d.errors.NotFound: pass + try: + await message.delete() + except d.errors.NotFound: + pass async def on_message(self, channel): def check(msg): - if msg.content == 'stop' and msg.channel is channel and msg.author.guild_permissions.administrator: + 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 @@ -153,14 +157,11 @@ class Administration: await self.queue.put(message) except exc.Abort: u.background['management']['auto_delete'].remove(channel.id) - u.update(u.background, 'background.json') + u.dump(u.background, 'background.pkl') print('Stopped looping {}'.format(channel.id)) await channel.send('✅ **Stopped deleting messages in** {}**.**'.format(channel.mention), delete_after=5) except AttributeError: pass - except Exception: - await channel.send(exc.base + '\n```' + traceback.format_exc(limit=1) + '```') - traceback.print_exc() @commands.command(name='autodelete', aliases=['autodel', 'ad']) @commands.has_permissions(administrator=True) @@ -171,13 +172,12 @@ class Administration: try: if channel.id not in u.background.setdefault('management', {}).setdefault('auto_delete', []): u.background['management']['auto_delete'].append(channel.id) - u.update(u.background, 'background.json') + u.dump(u.background, 'background.pkl') self.bot.loop.create_task(self.on_message(channel)) self.bot.loop.create_task(self.delete()) print('Looping {}'.format(channel.id)) await ctx.send('✅ **Auto-deleting all messages in this channel.**', 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) - except Exception: - await channel.send(exc.base + '\n```' + traceback.format_exc(limit=1) + '```') - traceback.print_exc() + else: + raise exc.Exists + except exc.Exists: + await ctx.send('❌ **Already auto-deleting in this channel.** Type `stop` to stop.', delete_after=10) diff --git a/src/main/cogs/owner.py b/src/main/cogs/owner.py index a0495a8..11fcd24 100644 --- a/src/main/cogs/owner.py +++ b/src/main/cogs/owner.py @@ -1,37 +1,93 @@ import asyncio import code import io -import pyrasite as pyr +import os import re import sys import traceback as tb import discord as d +import pyrasite as pyr from discord.ext import commands -from misc import checks from misc import exceptions as exc +from misc import checks +from utils import utils as u nl = re.compile('\n') + +class Bot: + + def __init__(self, bot, config): + self.bot = bot + self.config = config + + # Close connection to Discord - immediate offline + @commands.command(name=',die', aliases=[',d'], brief='Kills the bot', description='BOT OWNER ONLY\nCloses the connection to Discord', hidden=True) + @commands.is_owner() + @checks.del_ctx() + async def die(self, ctx): + if isinstance(self.bot.get_channel(self.config['startup_channel']), d.TextChannel): + await self.bot.get_channel(self.config['shutdown_channel']).send('**Shutting down...** 🌙') + # loop = self.bot.loop.all_tasks() + # for task in loop: + # task.cancel() + await u.session.close() + await self.bot.logout() + await self.bot.close() + print('-------') + print('CLOSED') + + @commands.command(name=',restart', aliases=[',res', ',r'], hidden=True) + @commands.is_owner() + @checks.del_ctx() + async def restart(self, ctx): + print('RESTARTING') + print('-------') + if isinstance(self.bot.get_channel(self.config['startup_channel']), d.TextChannel): + await self.bot.get_channel(self.config['shutdown_channel']).send('**Restarting...** 💤') + # loop = self.bot.loop.all_tasks() + # for task in loop: + # task.cancel() + await u.session.close() + await self.bot.logout() + await self.bot.close() + os.execl(sys.executable, 'python3', 'run.py') + + # Invite bot to bot owner's server + @commands.command(name=',invite', aliases=[',inv', ',link'], brief='Invite the bot', description='BOT OWNER ONLY\nInvite the bot to a server (Requires admin)', hidden=True) + @commands.is_owner() + @checks.del_ctx() + async def invite(self, ctx): + await ctx.send('🔗 https://discordapp.com/oauth2/authorize?&client_id={}&scope=bot&permissions={}'.format(self.config['client_id'], self.config['permissions']), delete_after=10) + + class Tools: def __init__(self, bot): self.bot = bot def format(self, i='', o=''): - if len(o) > 1: return '>>> {}\n{}'.format(i, o) - else: return '>>> {}'.format(i) + if len(o) > 1: + return '>>> {}\n{}'.format(i, o) + else: + return '>>> {}'.format(i) + async def generate(self, d, i='', o=''): return await d.send('```python\n{}```'.format(self.format(i, o))) + async def refresh(self, m, i='', o=''): global nl output = m.content[10:-3] - if len(nl.findall(output)) <= 20: await m.edit(content='```python\n{}\n{}\n>>>```'.format(output, self.format(i, o))) - else: await m.edit(content='```python\n{}```'.format(self.format(i, o))) + if len(nl.findall(output)) <= 20: + await m.edit(content='```python\n{}\n{}\n>>>```'.format(output, self.format(i, o))) + else: + await m.edit(content='```python\n{}```'.format(self.format(i, o))) async def generate_err(self, d, o=''): return await d.send('```\n{}```'.format(o)) + async def refresh_err(self, m, o=''): await m.edit(content='```\n{}```'.format(o)) @@ -42,8 +98,10 @@ class Tools: def execute(msg): if msg.content == ',exit' and msg.author is ctx.message.author: raise exc.CheckFail - elif msg.author is ctx.message.author and msg.channel is ctx.message.channel: return True - else: return False + elif msg.author is ctx.message.author and msg.channel is ctx.message.channel: + return True + else: + return False try: console = await self.generate(ctx) @@ -53,8 +111,10 @@ class Tools: await exe.delete() sys.stdout = io.StringIO() sys.stderr = io.StringIO() - try: exec(exe.content) - except Exception: tb.print_exc(limit=1) + try: + exec(exe.content) + except Exception: + tb.print_exc(limit=1) await self.refresh(console, exe.content, sys.stdout.getvalue()) await self.refresh_err(exception, sys.stderr.getvalue()) await ctx.send(console.content[10:-3]) @@ -62,9 +122,6 @@ class Tools: sys.stderr = sys.__stderr__ except exc.CheckFail: await ctx.send('↩️ **Exited console.**') - except Exception: - await ctx.send('{}\n```{}```'.format(exc.base, tb.format_exc(limit=1))) - tb.print_exc(limit=1, file=sys.__stderr__) finally: sys.stdout = sys.__stdout__ sys.stderr = sys.__stderr__ @@ -79,7 +136,7 @@ class Tools: exec(exe) await self.generate(ctx, exe, sys.stdout.getvalue()) except Exception: - await ctx.send('{}\n```{}```'.format(exc.base, tb.format_exc(limit=1))) + await ctx.send('```\n{}```'.format(tb.format_exc(limit=1))) tb.print_exc(limit=1) finally: sys.stdout = sys.__stdout__ @@ -90,9 +147,11 @@ class Tools: @checks.del_ctx() async def debug(self, ctx): console = await self.generate(ctx) + @debug.command(name='inject', aliases=['inj']) async def _inject(self, ctx, *, input_): pass + @debug.command(name='inspect', aliases=['ins']) async def _inspect(self, ctx, *, input_): pass diff --git a/src/main/cogs/tools.py b/src/main/cogs/tools.py index 875b88d..ace1ba5 100644 --- a/src/main/cogs/tools.py +++ b/src/main/cogs/tools.py @@ -1,29 +1,33 @@ import asyncio import datetime as dt -import discord -import httplib2 import mimetypes import os -import requests_oauthlib as ro import tempfile -import traceback +import traceback as tb import webbrowser + +import discord +import httplib2 +import requests_oauthlib as ro +from apiclient import http +from apiclient.discovery import build from discord.ext import commands +from oauth2client.client import flow_from_clientsecrets + #from run import config from cogs import booru -from misc import checks from misc import exceptions as exc +from misc import checks +from utils import utils as u from utils import formatter -from apiclient.discovery import build -from apiclient import http -from oauth2client.client import flow_from_clientsecrets youtube = None tempfile.tempdir = os.getcwd() command_dict = {} + class Utils: def __init__(self, bot): @@ -32,35 +36,25 @@ class Utils: @commands.command(name='last', aliases=['l', ','], brief='Reinvokes last command', description='Reinvokes previous command executed', hidden=True) async def last_command(self, ctx): global command_dict - try: - if command_dict.get(str(ctx.message.author.id), {}).get('args', None) is not None: - args = command_dict.get(str(ctx.message.author.id), {})['args'] - print(command_dict) - await ctx.invoke(command_dict.get(str(ctx.message.author.id), {}).get('command', None), args) - except Exception: - await ctx.send(exc.base + '\n```' + traceback.format_exc(limit=1) + '```') - traceback.print_exc(limit=1) + + if command_dict.get(str(ctx.message.author.id), {}).get('args', None) is not None: + args = command_dict.get(str(ctx.message.author.id), {})['args'] + print(command_dict) + await ctx.invoke(command_dict.get(str(ctx.message.author.id), {}).get('command', None), args) # [prefix]ping -> Pong! @commands.command(aliases=['p'], brief='Pong!', description='Returns latency from bot to Discord servers, not to user') @checks.del_ctx() async def ping(self, ctx): global command_dict - try: - await ctx.send(ctx.message.author.mention + ' 🏓 `' + str(int(self.bot.latency * 1000)) + 'ms`', delete_after=5) - except Exception: - await ctx.send(exc.base + '\n```' + traceback.format_exc(limit=1) + '```') - traceback.print_exc(limit=1) + + await ctx.send(ctx.message.author.mention + ' 🏓 `' + str(int(self.bot.latency * 1000)) + 'ms`', delete_after=5) command_dict.setdefault(str(ctx.message.author.id), {}).update({'command': ctx.command}) @commands.command(aliases=['pre'], brief='List bot prefixes', description='Shows all used prefixes') @checks.del_ctx() async def prefix(self, ctx): - try: - await ctx.send('**Prefix:** `,`') - except Exception: - await ctx.send(exc.base + '\n```' + traceback.format_exc(limit=1) + '```') - traceback.print_exc(limit=1) + await ctx.send('**Prefix:** `{}`'.format(u.config['prefix'])) @commands.group(name=',send', aliases=[',s'], hidden=True) @commands.is_owner() @@ -71,6 +65,7 @@ class Utils: @send.command(name='guild', aliases=['g', 'server', 's']) async def send_guild(self, ctx, guild, channel, *message): await discord.utils.get(self.bot.get_all_channels(), guild__name=guild, name=channel).send(formatter.tostring(message)) + @send.command(name='user', aliases=['u', 'member', 'm']) async def send_user(self, ctx, user, *message): await discord.utils.get(self.bot.get_all_members(), id=int(user)).send(formatter.tostring(message)) @@ -78,12 +73,14 @@ class Utils: @commands.command(aliases=['authenticateupload', 'authupload', 'authup', 'auth']) async def authenticate_upload(self, ctx): global youtube - flow = flow_from_clientsecrets('client_secrets.json', scope='https://www.googleapis.com/auth/youtube.upload', login_hint='botmyned@gmail.com', redirect_uri='urn:ietf:wg:oauth:2.0:oob') + flow = flow_from_clientsecrets('client_secrets.json', scope='https://www.googleapis.com/auth/youtube.upload', + login_hint='botmyned@gmail.com', redirect_uri='urn:ietf:wg:oauth:2.0:oob') flow.params['access_type'] = 'offline' webbrowser.open_new_tab(flow.step1_get_authorize_url()) credentials = flow.step2_exchange(input('Authorization code: ')) youtube = build('youtube', 'v3', http=credentials.authorize(http.build_http())) print('Service built.') + @commands.command(aliases=['up', 'u', 'vid', 'v']) @checks.is_listed() async def upload(self, ctx): @@ -100,16 +97,15 @@ class Utils: await attachments[0].save(temp) else: raise exc.InvalidVideoFile(mime) - print('https://www.youtube.com/watch?v=' + youtube.videos().insert(part='snippet', body={'categoryId': '24', 'title': 'Test'}, media_body=http.MediaFileUpload(temp.name, chunksize=-1))) + print('https://www.youtube.com/watch?v=' + youtube.videos().insert(part='snippet', + body={'categoryId': '24', 'title': 'Test'}, media_body=http.MediaFileUpload(temp.name, chunksize=-1))) except exc.InvalidVideoFile as e: await ctx.send('❌ `' + str(e) + '` **not valid video type.**', delete_after=10) except exc.TooManyAttachments as e: await ctx.send('❌ `' + str(e) + '` **too many attachments.** Only one attachment is permitted to upload.', delete_after=10) except exc.MissingAttachment: await ctx.send('❌ **Missing attachment.**', delete_after=10) - except Exception: - await ctx.send(exc.base + '\n```' + traceback.format_exc(limit=1) + '```') - traceback.print_exc(limit=1) + @upload.error async def upload_error(self, ctx, error): pass diff --git a/src/main/misc/checks.py b/src/main/misc/checks.py index 6876964..09daf75 100644 --- a/src/main/misc/checks.py +++ b/src/main/misc/checks.py @@ -1,7 +1,8 @@ import asyncio -import discord import json import traceback + +import discord from discord.ext import commands from discord.ext.commands import errors @@ -11,30 +12,43 @@ with open('config.json') as infile: owner_id = config['owner_id'] listed_ids = config['listed_ids'] + def is_owner(): async def predicate(ctx): return ctx.message.author.id == owner_id return commands.check(predicate) + + def is_admin(): def predicate(ctx): return ctx.message.author.guild_permissions.administrator return commands.check(predicate) + + def is_mod(): def predicate(ctx): return ctx.message.author.guild_permissions.ban_members return commands.check(predicate) + + def is_listed(): def predicate(ctx): return ctx.message.author.id in listed_ids return commands.check(predicate) + def owner(ctx): return ctx.message.author.id == owner_id + + def admin(ctx): return ctx.message.author.guild_permissions.administrator + + def mod(ctx): return ctx.message.author.guild_permissions.ban_members + def is_nsfw(): def predicate(ctx): if isinstance(ctx.message.channel, discord.TextChannel): @@ -42,6 +56,7 @@ def is_nsfw(): return True return commands.check(predicate) + def del_ctx(): async def predicate(ctx): if isinstance(ctx.message.channel, discord.TextChannel): diff --git a/src/main/misc/exceptions.py b/src/main/misc/exceptions.py index 8db2366..b772bd4 100644 --- a/src/main/misc/exceptions.py +++ b/src/main/misc/exceptions.py @@ -1,24 +1,93 @@ base = '⚠️ **An internal error has occurred.** Please notify my master! 🐺' -class Left(Exception): pass -class Right(Exception): pass -class Save(Exception): pass -class Exists(Exception): pass -class PostError(Exception): pass -class ImageError(Exception): pass -class MatchError(Exception): pass -class TagBlacklisted(Exception): pass -class BoundsError(Exception): pass -class TagBoundsError(Exception): pass -class TagExists(Exception): pass -class TagError(Exception): pass -class FlagError(Exception): pass -class BlacklistError(Exception): pass -class NotFound(Exception): pass -class Timeout(Exception): pass -class InvalidVideoFile(Exception): pass -class MissingAttachment(Exception): pass -class TooManyAttachments(Exception): pass -class CheckFail(Exception): pass -class Abort(Exception): pass -class Continue(Exception): pass + +class Left(Exception): + pass + + +class Right(Exception): + pass + + +class Save(Exception): + pass + + +class GoTo(Exception): + pass + + +class Exists(Exception): + pass + + +class PostError(Exception): + pass + + +class ImageError(Exception): + pass + + +class MatchError(Exception): + pass + + +class TagBlacklisted(Exception): + pass + + +class BoundsError(Exception): + pass + + +class TagBoundsError(Exception): + pass + + +class TagExists(Exception): + pass + + +class TagError(Exception): + pass + + +class FlagError(Exception): + pass + + +class BlacklistError(Exception): + pass + + +class NotFound(Exception): + pass + + +class Timeout(Exception): + pass + + +class InvalidVideoFile(Exception): + pass + + +class MissingAttachment(Exception): + pass + + +class TooManyAttachments(Exception): + pass + + +class CheckFail(Exception): + pass + + +class Abort(Exception): + pass + + +class Continue(Exception): + pass diff --git a/src/main/run.py b/src/main/run.py index abf32ad..73f8620 100644 --- a/src/main/run.py +++ b/src/main/run.py @@ -1,4 +1,21 @@ +import asyncio +import datetime as dt import json +import logging +import os +import subprocess +import sys +import traceback as tb + +import aiohttp as aio +import discord as d +from discord import utils +from discord.ext import commands + +from cogs import booru, info, management, owner, tools +from misc import exceptions as exc +from misc import checks +from utils import utils as u try: with open('config.json') as infile: @@ -6,24 +23,12 @@ try: print('\"config.json\" loaded.') except FileNotFoundError: with open('config.json', 'w') as outfile: - json.dump({'client_id': 0, 'listed_ids': [0], 'owner_id': 0, 'permissions': 126016, 'prefix': ',', 'shutdown_channel': 0, 'startup_channel': 0, 'token': 'str'}, outfile, indent=4, sort_keys=True) - raise FileNotFoundError('Config file not found: \"config.json\" created with abstract values. Restart \"run.py\" with correct values.') + json.dump({'client_id': 0, 'listed_ids': [0], 'owner_id': 0, 'permissions': 126016, 'prefix': ',', + 'shutdown_channel': 0, 'startup_channel': 0, 'token': 'str'}, outfile, indent=4, sort_keys=True) + raise FileNotFoundError( + 'Config file not found: \"config.json\" created with abstract values. Restart \"run.py\" with correct values.') -import asyncio -import datetime as dt -import discord as d -import os -import subprocess -import sys -import traceback -from discord import utils -from discord.ext import commands -from cogs import booru, info, owner, management, tools -from misc import checks -from misc import exceptions as exc -from utils import utils as u -import logging logging.basicConfig(level=logging.INFO) print('PID {}'.format(os.getpid())) @@ -31,16 +36,19 @@ print('PID {}'.format(os.getpid())) bot = commands.Bot(command_prefix=config['prefix'], description='Experimental booru bot') # Send and print ready message to #testing and console after logon + + @bot.event async def on_ready(): - global bot - bot.add_cog(tools.Utils(bot)) + bot.add_cog(owner.Bot(bot, config)) bot.add_cog(owner.Tools(bot)) bot.add_cog(management.Administration(bot)) bot.add_cog(info.Info(bot)) bot.add_cog(booru.MsG(bot)) + u.session = aio.ClientSession(loop=bot.loop) + # bot.loop.create_task(u.clear(booru.temp_urls, 30*60)) if isinstance(bot.get_channel(config['startup_channel']), d.TextChannel): @@ -49,71 +57,28 @@ async def on_ready(): print(bot.user.name) print('-------') -# Close connection to Discord - immediate offline -@bot.command(name=',die', aliases=[',d'], brief='Kills the bot', description='BOT OWNER ONLY\nCloses the connection to Discord', hidden=True) -@commands.is_owner() -@checks.del_ctx() -async def die(ctx): - try: - if isinstance(bot.get_channel(config['startup_channel']), d.TextChannel): - await bot.get_channel(config['shutdown_channel']).send('**Shutting down...** 🌙') - await bot.close() - print('-------') - print('CLOSED') - except Exception: - await ctx.send(exc.base + '\n```' + traceback.format_exc(limit=1) + '```') - traceback.print_exc(limit=1) -@bot.command(name=',restart', aliases=[',res', ',r'], hidden=True) -@commands.is_owner() -@checks.del_ctx() -async def restart(ctx): - try: - print('RESTARTING') - print('-------') - if isinstance(bot.get_channel(config['startup_channel']), d.TextChannel): - await bot.get_channel(config['shutdown_channel']).send('**Restarting...** 💤') - os.execl(sys.executable, 'python3', 'run.py') - except Exception: - await ctx.send('{}\n```{}```'.format(exc.base, traceback.format_exc(limit=1))) - traceback.print_exc(limit=1) +@bot.event +async def on_command_error(ctx, error): + print(error) + await ctx.send('{}\n```\n{}```'.format(exc.base, error)) -# Invite bot to bot owner's server -@bot.command(name=',invite', aliases=[',inv', ',link'], brief='Invite the bot', description='BOT OWNER ONLY\nInvite the bot to a server (Requires admin)', hidden=True) -@commands.is_owner() -@checks.del_ctx() -async def invite(ctx): - try: - await ctx.send('🔗 https://discordapp.com/oauth2/authorize?&client_id={}&scope=bot&permissions={}'.format(config['client_id'], config['permissions']), delete_after=10) - except Exception: - await ctx.send('{}\n```{}```'.format(exc.base, traceback.format_exc(limit=1))) - traceback.print_exc(limit=1) -@bot.command(brief='[IN TESTING]', description='[IN TESTING]', hidden=True) -async def hi(ctx): - user = ctx.message.author - try: - hello = 'Hewwo, {}.'.format(user.mention) - if user.id == checks.owner_id: - hello += '.. ***Master.*** uwu' - elif user.guild_permissions.administrator: - hello = '{} **Admin** {}'.format(hello[:7], hello[7:]) - elif user.guild_permissions.ban_members: - hello = '{} **Mod** {}'.format(hello[:7], hello[7:]) - await ctx.send(hello) - except Exception: - await ctx.send('{}\n```{}```'.format(exc.base, traceback.format_exc(limit=1))) - traceback.print_exc(limit=1) +async def on_reaction_add(r, u): + print('Reacted') + + +async def on_reaction_remove(r, u): + print('Removed') + @bot.command(name=',test', hidden=True) @commands.is_owner() @checks.del_ctx() async def test(ctx): - embed = d.Embed(title='/post/xxxxxx', url='https://static1.e621.net/data/4b/3e/4b3ec0c2e8580f418e4ce019dfd5ac32.png', timestamp=dt.datetime.utcnow(), color=ctx.me.color) - embed.set_image(url='https://static1.e621.net/data/27/0f/270fd28caa5e6d8bf542a76515848e02.png') - embed.set_footer(text='e621', icon_url='http://ndl.mgccw.com/mu3/app/20141013/18/1413204353554/icon/icon_xl.png') - embed.set_author(name='tags', url=ctx.message.author.avatar_url, icon_url=ctx.message.author.avatar_url) - embed.add_field(name='Link', value='https://static1.e621.net/data/c2/55/c255792b5a307ee6efa51d6bb3edf878.jpg') - await ctx.send(embed=embed) + test = await ctx.send('Test') + await test.add_reaction('✅') + bot.add_listener(on_reaction_add) + bot.add_listener(on_reaction_remove) bot.run(config['token']) diff --git a/src/main/utils/formatter.py b/src/main/utils/formatter.py index f8ed334..3dc87d4 100644 --- a/src/main/utils/formatter.py +++ b/src/main/utils/formatter.py @@ -10,6 +10,7 @@ def tostring(i, *, random=False): o = ' ' return o + def tostring_commas(i): if i: o = ',' @@ -18,6 +19,7 @@ def tostring_commas(i): return o[:-1] return '' + def dict_tostring(i): o = '' if i: @@ -25,6 +27,7 @@ def dict_tostring(i): o += '**' + k + ':** `' + tostring(v) + '`\n' return o + def dictelem_tostring(i): o = '' if i: diff --git a/src/main/utils/scraper.py b/src/main/utils/scraper.py index 6e3bb2e..818e41c 100644 --- a/src/main/utils/scraper.py +++ b/src/main/utils/scraper.py @@ -1,33 +1,17 @@ -import requests +import aiohttp as aio from bs4 import BeautifulSoup from lxml import html -from misc import exceptions as exc -def check_match(url): - r = requests.get(url) - soup = BeautifulSoup(r.content, 'html.parser') +from misc import exceptions as exc +from utils import utils as u + + +async def check_match(url): + r = await u.fetch('http://iqdb.harry.lu/?url={}'.format(url)) + soup = BeautifulSoup(await r.read(), 'html.parser') value = soup.find_all('a')[1].get('href') + if value != '#': return value else: raise exc.MatchError(value) - -def find_pool(url): - r = requests.get(url) - tree = html.fromstring(r.content) - post = tree.xpath('/html/body/div[@id="content"]/div[@id="pool-show"]/div[@style="margin-top: 2em;"]/span/a/@href') - print(post) - if post: - return post - else: - raise exc.PostError(post) - -def find_image_url(url): - r = requests.get(url) - tree = html.fromstring(r.content) - image_url = tree.xpath('/html/body/div[@id="content"]/div[@id="post-view"]/div[@class="content"]/div/img/@src') - print(image_url) - if image_url: - return image_url - else: - raise exc.ImageError(image_url) diff --git a/src/main/utils/utils.py b/src/main/utils/utils.py index eb007c8..50ec63a 100644 --- a/src/main/utils/utils.py +++ b/src/main/utils/utils.py @@ -1,26 +1,40 @@ +import asyncio import json +import pickle as pkl -try: - with open('background.json') as infile: - background = json.load(infile) - print('\"background.json\" loaded.') -except FileNotFoundError: - with open('background.json', 'w+') as iofile: - print('Background file not found: \"background.json\" created and loaded.') - json.dump({}, iofile, indent=4, sort_keys=True) - iofile.seek(0) - background = json.load(iofile) +import aiohttp as aio + + +def setdefault(filename, default=None): + try: + with open(filename, 'rb') as infile: + print('\"{}\" loaded.'.format(filename)) + return pkl.load(infile) + except FileNotFoundError: + with open(filename, 'wb+') as iofile: + print('File not found: \"{}\" created and loaded with default values.'.format(filename)) + pkl.dump(default, iofile) + iofile.seek(0) + return pkl.load(iofile) + + +def load(filename): + with open(filename, 'rb') as infile: + return pkl.load(infile) + + +def dump(obj, filename): + with open(filename, 'wb') as outfile: + pkl.dump(obj, outfile) + + +background = setdefault('./cogs/background.pkl', {}) with open('config.json') as infile: config = json.load(infile) -def update(out, file): - with open(file, 'w') as outfile: - json.dump(out, outfile, indent=4, sort_keys=True) -import asyncio - -async def clear(obj, interval=10*60, replace=None): +async def clear(obj, interval=10 * 60, replace=None): if replace is None: if type(obj) is list: replace = [] @@ -34,3 +48,16 @@ async def clear(obj, interval=10*60, replace=None): while True: obj = replace asyncio.sleep(interval) + + +session = None + +HEADERS = {'user-agent': 'Modumind/0.0.1 (Myned)'} + + +async def fetch(url, *, params={}, json=False): + global session, HEADERS + async with session.get(url, params=params, headers=HEADERS) as r: + if json is True: + return await r.json() + return r