From c8caeadec43b63956681ba76fab5b27e3f9f8686 Mon Sep 17 00:00:00 2001 From: Ske Date: Sat, 27 Oct 2018 21:03:35 +0200 Subject: [PATCH 1/9] Change discord.py dependency to rewrite branch pinned at 860d6a9a --- src/requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/requirements.txt b/src/requirements.txt index c08261af..95007db4 100644 --- a/src/requirements.txt +++ b/src/requirements.txt @@ -2,7 +2,7 @@ aiodns aiohttp asyncpg dateparser -discord.py +https://github.com/Rapptz/discord.py/archive/860d6a9ace8248dfeec18b8b159e7b757d9f56bb.zip#egg=discord.py humanize uvloop ciso8601 From 4217d5d5d894a8f021402be8ae1a5c889b6fcad9 Mon Sep 17 00:00:00 2001 From: Ske Date: Sat, 27 Oct 2018 22:00:41 +0200 Subject: [PATCH 2/9] Basic proxy functionality fixed --- src/bot_main.py | 5 +- src/pluralkit/bot/__init__.py | 264 +++++++----- src/pluralkit/bot/commands/__init__.py | 2 +- src/pluralkit/bot/proxy.py | 532 ++++++++++++++----------- src/pluralkit/bot/utils.py | 11 +- src/pluralkit/db.py | 54 +-- 6 files changed, 499 insertions(+), 369 deletions(-) diff --git a/src/bot_main.py b/src/bot_main.py index 3d592a6e..bb8886b6 100644 --- a/src/bot_main.py +++ b/src/bot_main.py @@ -5,7 +5,4 @@ import uvloop asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) from pluralkit import bot - -pk = bot.PluralKitBot(os.environ["TOKEN"]) -loop = asyncio.get_event_loop() -loop.run_until_complete(pk.run()) \ No newline at end of file +bot.run() \ No newline at end of file diff --git a/src/pluralkit/bot/__init__.py b/src/pluralkit/bot/__init__.py index 3a432e83..1d2708bc 100644 --- a/src/pluralkit/bot/__init__.py +++ b/src/pluralkit/bot/__init__.py @@ -1,127 +1,187 @@ -import asyncio -import json -import logging -import os -import time +import asyncpg +import sys -import traceback -from datetime import datetime +import asyncio +import os + +import logging import discord from pluralkit import db -from pluralkit.bot import channel_logger, commands, proxy, embeds +from pluralkit.bot import commands, proxy logging.basicConfig(level=logging.INFO, format="[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s") +def connect_to_database() -> asyncpg.pool.Pool: + username = os.environ["DATABASE_USER"] + password = os.environ["DATABASE_PASS"] + name = os.environ["DATABASE_NAME"] + host = os.environ["DATABASE_HOST"] + port = os.environ["DATABASE_PORT"] -# logging.getLogger("pluralkit").setLevel(logging.DEBUG) + if username is None or password is None or name is None or host is None or port is None: + print("Database credentials not specified. Please pass valid PostgreSQL database credentials in the DATABASE_[USER|PASS|NAME|HOST|PORT] environment variable.", file=sys.stderr) + sys.exit(1) -class PluralKitBot: - def __init__(self, token): - self.token = token - self.logger = logging.getLogger("pluralkit.bot") + try: + port = int(port) + except ValueError: + print("Please pass a valid integer as the DATABASE_PORT environment variable.", file=sys.stderr) + sys.exit(1) - self.client = discord.Client() - self.client.event(self.on_error) - self.client.event(self.on_ready) - self.client.event(self.on_message) - self.client.event(self.on_socket_raw_receive) + return asyncio.get_event_loop().run_until_complete(db.connect( + username=username, + password=password, + database=name, + host=host, + port=port + )) - self.channel_logger = channel_logger.ChannelLogger(self.client) +def run(): + pool = connect_to_database() - self.proxy = proxy.Proxy(self.client, token, self.channel_logger) + client = discord.Client() - async def on_error(self, evt, *args, **kwargs): - self.logger.exception("Error while handling event {} with arguments {}:".format(evt, args)) + @client.event + async def on_ready(): + print("PluralKit started.") + print("User: {}#{} (ID: {})".format(client.user.name, client.user.discriminator, client.user.id)) + print("{} servers".format(len(client.guilds))) + print("{} shards".format(client.shard_count or 1)) - async def on_ready(self): - self.logger.info("Connected to Discord.") - self.logger.info("- Account: {}#{}".format(self.client.user.name, self.client.user.discriminator)) - self.logger.info("- User ID: {}".format(self.client.user.id)) - self.logger.info("- {} servers".format(len(self.client.servers))) - - async def on_message(self, message): - # Ignore bot messages + @client.event + async def on_message(message: discord.Message): + # Ignore messages from bots if message.author.bot: return - try: - if await self.handle_command_dispatch(message): + # Grab a database connection from the pool + async with pool.acquire() as conn: + # First pass: do command handling + did_run_command = await commands.command_dispatch(client, message, conn) + if did_run_command: return - if await self.handle_proxy_dispatch(message): - return - except Exception: - await self.log_error_in_channel(message) + # Second pass: do proxy matching + await proxy.try_proxy_message(message, conn) - async def on_socket_raw_receive(self, msg): - # Since on_reaction_add is buggy (only works for messages the bot's already cached, ie. no old messages) - # we parse socket data manually for the reaction add event - if isinstance(msg, str): - try: - msg_data = json.loads(msg) - if msg_data.get("t") == "MESSAGE_REACTION_ADD": - evt_data = msg_data.get("d") - if evt_data: - user_id = evt_data["user_id"] - message_id = evt_data["message_id"] - emoji = evt_data["emoji"]["name"] - async with self.pool.acquire() as conn: - await self.proxy.handle_reaction(conn, user_id, message_id, emoji) - elif msg_data.get("t") == "MESSAGE_DELETE": - evt_data = msg_data.get("d") - if evt_data: - message_id = evt_data["id"] - async with self.pool.acquire() as conn: - await self.proxy.handle_deletion(conn, message_id) - except ValueError: - pass + bot_token = os.environ["TOKEN"] + if not bot_token: + print("No token specified. Please pass a valid Discord bot token in the TOKEN environment variable.", + file=sys.stderr) + sys.exit(1) - async def handle_command_dispatch(self, message): - async with self.pool.acquire() as conn: - result = await commands.command_dispatch(self.client, message, conn) - return result + client.run(bot_token) - async def handle_proxy_dispatch(self, message): - # Try doing proxy parsing - async with self.pool.acquire() as conn: - return await self.proxy.try_proxy_message(conn, message) +# logging.getLogger("pluralkit").setLevel(logging.DEBUG) - async def log_error_in_channel(self, message): - channel_id = os.environ["LOG_CHANNEL"] - if not channel_id: - return - - channel = self.client.get_channel(channel_id) - embed = embeds.exception_log( - message.content, - message.author.name, - message.author.discriminator, - message.server.id if message.server else None, - message.channel.id - ) - - await self.client.send_message(channel, "```python\n{}```".format(traceback.format_exc()), embed=embed) - - async def run(self): - try: - self.logger.info("Connecting to database...") - self.pool = await db.connect( - os.environ["DATABASE_USER"], - os.environ["DATABASE_PASS"], - os.environ["DATABASE_NAME"], - os.environ["DATABASE_HOST"], - int(os.environ["DATABASE_PORT"]) - ) - - self.logger.info("Attempting to create tables...") - async with self.pool.acquire() as conn: - await db.create_tables(conn) - - self.logger.info("Connecting to Discord...") - await self.client.start(self.token) - finally: - self.logger.info("Logging out from Discord...") - await self.client.logout() +# class PluralKitBot: +# def __init__(self, token): +# self.token = token +# self.logger = logging.getLogger("pluralkit.bot") +# +# self.client = discord.Client() +# self.client.event(self.on_error) +# self.client.event(self.on_ready) +# self.client.event(self.on_message) +# self.client.event(self.on_socket_raw_receive) +# +# self.channel_logger = channel_logger.ChannelLogger(self.client) +# +# self.proxy = proxy.Proxy(self.client, token, self.channel_logger) +# +# async def on_error(self, evt, *args, **kwargs): +# self.logger.exception("Error while handling event {} with arguments {}:".format(evt, args)) +# +# async def on_ready(self): +# self.logger.info("Connected to Discord.") +# self.logger.info("- Account: {}#{}".format(self.client.user.name, self.client.user.discriminator)) +# self.logger.info("- User ID: {}".format(self.client.user.id)) +# self.logger.info("- {} servers".format(len(self.client.servers))) +# +# async def on_message(self, message): +# # Ignore bot messages +# if message.author.bot: +# return +# +# try: +# if await self.handle_command_dispatch(message): +# return +# +# if await self.handle_proxy_dispatch(message): +# return +# except Exception: +# await self.log_error_in_channel(message) +# +# async def on_socket_raw_receive(self, msg): +# # Since on_reaction_add is buggy (only works for messages the bot's already cached, ie. no old messages) +# # we parse socket data manually for the reaction add event +# if isinstance(msg, str): +# try: +# msg_data = json.loads(msg) +# if msg_data.get("t") == "MESSAGE_REACTION_ADD": +# evt_data = msg_data.get("d") +# if evt_data: +# user_id = evt_data["user_id"] +# message_id = evt_data["message_id"] +# emoji = evt_data["emoji"]["name"] +# +# async with self.pool.acquire() as conn: +# await self.proxy.handle_reaction(conn, user_id, message_id, emoji) +# elif msg_data.get("t") == "MESSAGE_DELETE": +# evt_data = msg_data.get("d") +# if evt_data: +# message_id = evt_data["id"] +# async with self.pool.acquire() as conn: +# await self.proxy.handle_deletion(conn, message_id) +# except ValueError: +# pass +# +# async def handle_command_dispatch(self, message): +# async with self.pool.acquire() as conn: +# result = await commands.command_dispatch(self.client, message, conn) +# return result +# +# async def handle_proxy_dispatch(self, message): +# # Try doing proxy parsing +# async with self.pool.acquire() as conn: +# return await self.proxy.try_proxy_message(conn, message) +# +# async def log_error_in_channel(self, message): +# channel_id = os.environ["LOG_CHANNEL"] +# if not channel_id: +# return +# +# channel = self.client.get_channel(channel_id) +# embed = embeds.exception_log( +# message.content, +# message.author.name, +# message.author.discriminator, +# message.server.id if message.server else None, +# message.channel.id +# ) +# +# await self.client.send_message(channel, "```python\n{}```".format(traceback.format_exc()), embed=embed) +# +# async def run(self): +# try: +# self.logger.info("Connecting to database...") +# self.pool = await db.connect( +# os.environ["DATABASE_USER"], +# os.environ["DATABASE_PASS"], +# os.environ["DATABASE_NAME"], +# os.environ["DATABASE_HOST"], +# int(os.environ["DATABASE_PORT"]) +# ) +# +# self.logger.info("Attempting to create tables...") +# async with self.pool.acquire() as conn: +# await db.create_tables(conn) +# +# self.logger.info("Connecting to Discord...") +# await self.client.start(self.token) +# finally: +# self.logger.info("Logging out from Discord...") +# await self.client.logout() diff --git a/src/pluralkit/bot/commands/__init__.py b/src/pluralkit/bot/commands/__init__.py index 565ba6d4..6f346b28 100644 --- a/src/pluralkit/bot/commands/__init__.py +++ b/src/pluralkit/bot/commands/__init__.py @@ -118,7 +118,7 @@ class CommandContext: raise CommandError("Timed out - try again.") return reaction.reaction.emoji == "✅" - async def confirm_text(self, user: discord.Member, channel: discord.Channel, confirm_text: str, message: str): + async def confirm_text(self, user: discord.Member, channel: discord.TextChannel, confirm_text: str, message: str): await self.reply(message) message = await self.client.wait_for_message(channel=channel, author=user, timeout=60.0*5) diff --git a/src/pluralkit/bot/proxy.py b/src/pluralkit/bot/proxy.py index 2d07c5b8..a6920531 100644 --- a/src/pluralkit/bot/proxy.py +++ b/src/pluralkit/bot/proxy.py @@ -1,17 +1,16 @@ -import ciso8601 +from io import BytesIO + +import discord import logging import re -import time -from typing import List, Optional - -import aiohttp -import discord +from typing import List from pluralkit import db -from pluralkit.bot import channel_logger, utils, embeds +from pluralkit.bot import utils logger = logging.getLogger("pluralkit.bot.proxy") + def extract_leading_mentions(message_text): # This regex matches one or more mentions at the start of a message, separated by any amount of spaces match = re.match(r"^(<(@|@!|#|@&|a?:\w+:)\d+>\s*)+", message_text) @@ -34,7 +33,9 @@ def match_member_proxy_tags(member: db.ProxyMember, message_text: str): # Ignore mentions at the very start of the message, and match proxy tags after those message_text, leading_mentions = extract_leading_mentions(message_text) - logger.debug("Matching text '{}' and leading mentions '{}' to proxy tags {}text{}".format(message_text, leading_mentions, prefix, suffix)) + logger.debug( + "Matching text '{}' and leading mentions '{}' to proxy tags {}text{}".format(message_text, leading_mentions, + prefix, suffix)) if message_text.startswith(member.prefix or "") and message_text.endswith(member.suffix or ""): prefix_length = len(prefix) @@ -59,243 +60,314 @@ def match_proxy_tags(members: List[db.ProxyMember], message_text: str): for member in members: match = match_member_proxy_tags(member, message_text) - if match is not None: # Using "is not None" because an empty string is OK here too + if match is not None: # Using "is not None" because an empty string is OK here too logger.debug("Matched member {} with inner text '{}'".format(member.hid, match)) return member, match -def get_message_attachment_url(message: discord.Message): +async def get_or_create_webhook_for_channel(conn, channel: discord.TextChannel): + # First, check if we have one saved in the DB + webhook_from_db = await db.get_webhook(conn, channel.id) + if webhook_from_db: + webhook_id, webhook_token = webhook_from_db + + session = channel._state.http._session + hook = discord.Webhook.partial(webhook_id, webhook_token, adapter=discord.AsyncWebhookAdapter(session)) + + # Workaround for https://github.com/Rapptz/discord.py/issues/1242 + hook._adapter.store_user = hook._adapter._store_user + return hook + + # If not, we create one and save it + created_webhook = await channel.create_webhook(name="PluralKit Proxy Webhook") + await db.add_webhook(conn, channel.id, created_webhook.id, created_webhook.token) + return created_webhook + + +async def make_attachment_file(message: discord.Message): if not message.attachments: return None - attachment = message.attachments[0] - if "proxy_url" in attachment: - return attachment["proxy_url"] + first_attachment = message.attachments[0] - if "url" in attachment: - return attachment["url"] + # Copy the file data to the buffer + # TODO: do this without buffering... somehow + bio = BytesIO() + await first_attachment.save(bio) + + return discord.File(bio, first_attachment.filename) -# TODO: possibly move this to bot __init__ so commands can access it too -class WebhookPermissionError(Exception): - pass +async def do_proxy_message(conn, original_message: discord.Message, proxy_member: db.ProxyMember, + inner_text: str): + # Send the message through the webhook + webhook = await get_or_create_webhook_for_channel(conn, original_message.channel) + sent_message = await webhook.send( + content=inner_text, + username=proxy_member.name, + avatar_url=proxy_member.avatar_url, + file=await make_attachment_file(original_message), + wait=True + ) + + # Save the proxied message in the database + await db.add_message(conn, sent_message.id, original_message.channel.id, proxy_member.id, + original_message.author.id) + + # TODO: log message in log channel -class DeletionPermissionError(Exception): - pass +async def try_proxy_message(message: discord.Message, conn) -> bool: + # Don't bother proxying in DMs with the bot + if isinstance(message.channel, discord.abc.PrivateChannel): + return False + # Return every member associated with the account + members = await db.get_members_by_account(conn, message.author.id) + proxy_match = match_proxy_tags(members, message.content) + if not proxy_match: + # No proxy tags match here, we done + return False -class Proxy: - def __init__(self, client: discord.Client, token: str, logger: channel_logger.ChannelLogger): - self.logger = logging.getLogger("pluralkit.bot.proxy") - self.session = aiohttp.ClientSession() - self.client = client - self.token = token - self.channel_logger = logger + member, inner_text = proxy_match - async def save_channel_webhook(self, conn, channel: discord.Channel, id: str, token: str) -> (str, str): - await db.add_webhook(conn, channel.id, id, token) - return id, token + # Sanitize inner text for @everyones and such + inner_text = utils.sanitize(inner_text) - async def create_and_add_channel_webhook(self, conn, channel: discord.Channel) -> (str, str): - # This method is only called if there's no webhook found in the DB (and hopefully within a transaction) - # No need to worry about error handling if there's a DB conflict (which will throw an exception because DB constraints) - req_headers = {"Authorization": "Bot {}".format(self.token)} + # If we don't have an inner text OR an attachment, we cancel because the hook can't send that + if not inner_text and not message.attachments: + return False - # First, check if there's already a webhook belonging to the bot - async with self.session.get("https://discordapp.com/api/v6/channels/{}/webhooks".format(channel.id), - headers=req_headers) as resp: - if resp.status == 200: - webhooks = await resp.json() - for webhook in webhooks: - if webhook["user"]["id"] == self.client.user.id: - # This webhook belongs to us, we can use that, return it and save it - return await self.save_channel_webhook(conn, channel, webhook["id"], webhook["token"]) - elif resp.status == 403: - self.logger.warning( - "Did not have permission to fetch webhook list (server={}, channel={})".format(channel.server.id, - channel.id)) - raise WebhookPermissionError() - else: - raise discord.HTTPException(resp, await resp.text()) + # So, we now have enough information to successfully proxy a message + async with conn.transaction(): + await do_proxy_message(conn, message, member, inner_text) + return True - # Then, try submitting a new one - req_data = {"name": "PluralKit Proxy Webhook"} - async with self.session.post("https://discordapp.com/api/v6/channels/{}/webhooks".format(channel.id), - json=req_data, headers=req_headers) as resp: - if resp.status == 200: - webhook = await resp.json() - return await self.save_channel_webhook(conn, channel, webhook["id"], webhook["token"]) - elif resp.status == 403: - self.logger.warning( - "Did not have permission to create webhook (server={}, channel={})".format(channel.server.id, - channel.id)) - raise WebhookPermissionError() - else: - raise discord.HTTPException(resp, await resp.text()) - - # Should not be reached without an exception being thrown - - async def get_webhook_for_channel(self, conn, channel: discord.Channel): - async with conn.transaction(): - hook_match = await db.get_webhook(conn, channel.id) - if not hook_match: - # We don't have a webhook, create/add one - return await self.create_and_add_channel_webhook(conn, channel) - else: - return hook_match - - async def do_proxy_message(self, conn, member: db.ProxyMember, original_message: discord.Message, text: str, - attachment_url: str, has_already_retried=False): - hook_id, hook_token = await self.get_webhook_for_channel(conn, original_message.channel) - - form_data = aiohttp.FormData() - form_data.add_field("username", "{} {}".format(member.name, member.tag or "").strip()) - - if text: - form_data.add_field("content", text) - - if attachment_url: - attachment_resp = await self.session.get(attachment_url) - form_data.add_field("file", attachment_resp.content, content_type=attachment_resp.content_type, - filename=attachment_resp.url.name) - - if member.avatar_url: - form_data.add_field("avatar_url", member.avatar_url) - - async with self.session.post( - "https://discordapp.com/api/v6/webhooks/{}/{}?wait=true".format(hook_id, hook_token), - data=form_data) as resp: - if resp.status == 200: - message = await resp.json() - - await db.add_message(conn, message["id"], message["channel_id"], member.id, original_message.author.id) - - try: - await self.client.delete_message(original_message) - except discord.Forbidden: - self.logger.warning( - "Did not have permission to delete original message (server={}, channel={})".format( - original_message.server.id, original_message.channel.id)) - raise DeletionPermissionError() - except discord.NotFound: - self.logger.warning("Tried to delete message when proxying, but message was already gone (server={}, channel={})".format(original_message.server.id, original_message.channel.id)) - - message_image = None - if message["attachments"]: - first_attachment = message["attachments"][0] - if "width" in first_attachment and "height" in first_attachment: - # Only log attachments that are actually images - message_image = first_attachment["url"] - - await self.channel_logger.log_message_proxied(conn, - server_id=original_message.server.id, - channel_name=original_message.channel.name, - channel_id=original_message.channel.id, - sender_name=original_message.author.name, - sender_disc=original_message.author.discriminator, - sender_id=original_message.author.id, - member_name=member.name, - member_hid=member.hid, - member_avatar_url=member.avatar_url, - system_name=member.system_name, - system_hid=member.system_hid, - message_text=text, - message_image=message_image, - message_timestamp=ciso8601.parse_datetime( - message["timestamp"]), - message_id=message["id"]) - elif resp.status == 404 and not has_already_retried: - # Webhook doesn't exist. Delete it from the DB, create, and add a new one - self.logger.warning("Webhook registered in DB doesn't exist, deleting hook from DB, re-adding, and trying again (channel={}, hook={})".format(original_message.channel.id, hook_id)) - await db.delete_webhook(conn, original_message.channel.id) - await self.create_and_add_channel_webhook(conn, original_message.channel) - - # Then try again all over, making sure to not retry again and go in a loop should it continually fail - return await self.do_proxy_message(conn, member, original_message, text, attachment_url, has_already_retried=True) - else: - raise discord.HTTPException(resp, await resp.text()) - - async def try_proxy_message(self, conn, message: discord.Message): - # Can't proxy in DMs, webhook creation will explode - if message.channel.is_private: - return False - - # Big fat query to find every member associated with this account - # Returned member object has a few more keys (system tag, for example) - members = await db.get_members_by_account(conn, account_id=message.author.id) - - match = match_proxy_tags(members, message.content) - if not match: - return False - - member, text = match - attachment_url = get_message_attachment_url(message) - - # Can't proxy a message with no text AND no attachment - if not text and not attachment_url: - self.logger.debug("Skipping message because of no text and no attachment") - return False - - # Remember to sanitize the text (remove @everyones and such) - text = utils.sanitize(text) - - try: - async with conn.transaction(): - await self.do_proxy_message(conn, member, message, text=text, attachment_url=attachment_url) - except WebhookPermissionError: - embed = embeds.error("PluralKit does not have permission to manage webhooks for this channel. Contact your local server administrator to fix this.") - await self.client.send_message(message.channel, embed=embed) - except DeletionPermissionError: - embed = embeds.error("PluralKit does not have permission to delete messages in this channel. Contact your local server administrator to fix this.") - await self.client.send_message(message.channel, embed=embed) - - return True - - async def try_delete_message(self, conn, message_id: str, check_user_id: Optional[str], delete_message: bool, deleted_by_moderator: bool): - async with conn.transaction(): - # Find the message in the DB, and make sure it's sent by the user (if we need to check) - if check_user_id: - db_message = await db.get_message_by_sender_and_id(conn, message_id=message_id, sender_id=check_user_id) - else: - db_message = await db.get_message(conn, message_id=message_id) - - if db_message: - self.logger.debug("Deleting message {}".format(message_id)) - channel = self.client.get_channel(str(db_message.channel)) - - # If we should also delete the actual message, do that - if delete_message: - message = await self.client.get_message(channel, message_id) - - try: - await self.client.delete_message(message) - except discord.Forbidden: - self.logger.warning( - "Did not have permission to remove message, aborting deletion (server={}, channel={})".format( - channel.server.id, channel.id)) - return - - # Remove it from the DB - await db.delete_message(conn, message_id) - - # Then log deletion to logging channel - await self.channel_logger.log_message_deleted(conn, - server_id=channel.server.id, - channel_name=channel.name, - member_name=db_message.name, - member_hid=db_message.hid, - member_avatar_url=db_message.avatar_url, - system_name=db_message.system_name, - system_hid=db_message.system_hid, - message_text=db_message.content, - message_id=message_id, - deleted_by_moderator=deleted_by_moderator) - - async def handle_reaction(self, conn, user_id: str, message_id: str, emoji: str): - if emoji == "❌": - await self.try_delete_message(conn, message_id, check_user_id=user_id, delete_message=True, deleted_by_moderator=False) - - async def handle_deletion(self, conn, message_id: str): - # Don't delete the message, it's already gone at this point, just handle DB deletion and logging - await self.try_delete_message(conn, message_id, check_user_id=None, delete_message=False, deleted_by_moderator=True) +# # TODO: possibly move this to bot __init__ so commands can access it too +# class WebhookPermissionError(Exception): +# pass +# +# +# class DeletionPermissionError(Exception): +# pass +# +# +# class Proxy: +# def __init__(self, client: discord.Client, token: str, logger: channel_logger.ChannelLogger): +# self.logger = logging.getLogger("pluralkit.bot.proxy") +# self.session = aiohttp.ClientSession() +# self.client = client +# self.token = token +# self.channel_logger = logger +# +# async def save_channel_webhook(self, conn, channel: discord.Channel, id: str, token: str) -> (str, str): +# await db.add_webhook(conn, channel.id, id, token) +# return id, token +# +# async def create_and_add_channel_webhook(self, conn, channel: discord.Channel) -> (str, str): +# # This method is only called if there's no webhook found in the DB (and hopefully within a transaction) +# # No need to worry about error handling if there's a DB conflict (which will throw an exception because DB constraints) +# req_headers = {"Authorization": "Bot {}".format(self.token)} +# +# # First, check if there's already a webhook belonging to the bot +# async with self.session.get("https://discordapp.com/api/v6/channels/{}/webhooks".format(channel.id), +# headers=req_headers) as resp: +# if resp.status == 200: +# webhooks = await resp.json() +# for webhook in webhooks: +# if webhook["user"]["id"] == self.client.user.id: +# # This webhook belongs to us, we can use that, return it and save it +# return await self.save_channel_webhook(conn, channel, webhook["id"], webhook["token"]) +# elif resp.status == 403: +# self.logger.warning( +# "Did not have permission to fetch webhook list (server={}, channel={})".format(channel.server.id, +# channel.id)) +# raise WebhookPermissionError() +# else: +# raise discord.HTTPException(resp, await resp.text()) +# +# # Then, try submitting a new one +# req_data = {"name": "PluralKit Proxy Webhook"} +# async with self.session.post("https://discordapp.com/api/v6/channels/{}/webhooks".format(channel.id), +# json=req_data, headers=req_headers) as resp: +# if resp.status == 200: +# webhook = await resp.json() +# return await self.save_channel_webhook(conn, channel, webhook["id"], webhook["token"]) +# elif resp.status == 403: +# self.logger.warning( +# "Did not have permission to create webhook (server={}, channel={})".format(channel.server.id, +# channel.id)) +# raise WebhookPermissionError() +# else: +# raise discord.HTTPException(resp, await resp.text()) +# +# # Should not be reached without an exception being thrown +# +# async def get_webhook_for_channel(self, conn, channel: discord.Channel): +# async with conn.transaction(): +# hook_match = await db.get_webhook(conn, channel.id) +# if not hook_match: +# # We don't have a webhook, create/add one +# return await self.create_and_add_channel_webhook(conn, channel) +# else: +# return hook_match +# +# async def do_proxy_message(self, conn, member: db.ProxyMember, original_message: discord.Message, text: str, +# attachment_url: str, has_already_retried=False): +# hook_id, hook_token = await self.get_webhook_for_channel(conn, original_message.channel) +# +# form_data = aiohttp.FormData() +# form_data.add_field("username", "{} {}".format(member.name, member.tag or "").strip()) +# +# if text: +# form_data.add_field("content", text) +# +# if attachment_url: +# attachment_resp = await self.session.get(attachment_url) +# form_data.add_field("file", attachment_resp.content, content_type=attachment_resp.content_type, +# filename=attachment_resp.url.name) +# +# if member.avatar_url: +# form_data.add_field("avatar_url", member.avatar_url) +# +# async with self.session.post( +# "https://discordapp.com/api/v6/webhooks/{}/{}?wait=true".format(hook_id, hook_token), +# data=form_data) as resp: +# if resp.status == 200: +# message = await resp.json() +# +# await db.add_message(conn, message["id"], message["channel_id"], member.id, original_message.author.id) +# +# try: +# await self.client.delete_message(original_message) +# except discord.Forbidden: +# self.logger.warning( +# "Did not have permission to delete original message (server={}, channel={})".format( +# original_message.server.id, original_message.channel.id)) +# raise DeletionPermissionError() +# except discord.NotFound: +# self.logger.warning("Tried to delete message when proxying, but message was already gone (server={}, channel={})".format(original_message.server.id, original_message.channel.id)) +# +# message_image = None +# if message["attachments"]: +# first_attachment = message["attachments"][0] +# if "width" in first_attachment and "height" in first_attachment: +# # Only log attachments that are actually images +# message_image = first_attachment["url"] +# +# await self.channel_logger.log_message_proxied(conn, +# server_id=original_message.server.id, +# channel_name=original_message.channel.name, +# channel_id=original_message.channel.id, +# sender_name=original_message.author.name, +# sender_disc=original_message.author.discriminator, +# sender_id=original_message.author.id, +# member_name=member.name, +# member_hid=member.hid, +# member_avatar_url=member.avatar_url, +# system_name=member.system_name, +# system_hid=member.system_hid, +# message_text=text, +# message_image=message_image, +# message_timestamp=ciso8601.parse_datetime( +# message["timestamp"]), +# message_id=message["id"]) +# elif resp.status == 404 and not has_already_retried: +# # Webhook doesn't exist. Delete it from the DB, create, and add a new one +# self.logger.warning("Webhook registered in DB doesn't exist, deleting hook from DB, re-adding, and trying again (channel={}, hook={})".format(original_message.channel.id, hook_id)) +# await db.delete_webhook(conn, original_message.channel.id) +# await self.create_and_add_channel_webhook(conn, original_message.channel) +# +# # Then try again all over, making sure to not retry again and go in a loop should it continually fail +# return await self.do_proxy_message(conn, member, original_message, text, attachment_url, has_already_retried=True) +# else: +# raise discord.HTTPException(resp, await resp.text()) +# +# async def try_proxy_message(self, conn, message: discord.Message): +# # Can't proxy in DMs, webhook creation will explode +# if message.channel.is_private: +# return False +# +# # Big fat query to find every member associated with this account +# # Returned member object has a few more keys (system tag, for example) +# members = await db.get_members_by_account(conn, account_id=message.author.id) +# +# match = match_proxy_tags(members, message.content) +# if not match: +# return False +# +# member, text = match +# attachment_url = get_message_attachment_url(message) +# +# # Can't proxy a message with no text AND no attachment +# if not text and not attachment_url: +# self.logger.debug("Skipping message because of no text and no attachment") +# return False +# +# # Remember to sanitize the text (remove @everyones and such) +# text = utils.sanitize(text) +# +# try: +# async with conn.transaction(): +# await self.do_proxy_message(conn, member, message, text=text, attachment_url=attachment_url) +# except WebhookPermissionError: +# embed = embeds.error("PluralKit does not have permission to manage webhooks for this channel. Contact your local server administrator to fix this.") +# await self.client.send_message(message.channel, embed=embed) +# except DeletionPermissionError: +# embed = embeds.error("PluralKit does not have permission to delete messages in this channel. Contact your local server administrator to fix this.") +# await self.client.send_message(message.channel, embed=embed) +# +# return True +# +# async def try_delete_message(self, conn, message_id: str, check_user_id: Optional[str], delete_message: bool, deleted_by_moderator: bool): +# async with conn.transaction(): +# # Find the message in the DB, and make sure it's sent by the user (if we need to check) +# if check_user_id: +# db_message = await db.get_message_by_sender_and_id(conn, message_id=message_id, sender_id=check_user_id) +# else: +# db_message = await db.get_message(conn, message_id=message_id) +# +# if db_message: +# self.logger.debug("Deleting message {}".format(message_id)) +# channel = self.client.get_channel(str(db_message.channel)) +# +# # Fetch the original message from the server +# try: +# original_message = await self.client.get_message(channel, message_id) +# except discord.NotFound: +# # Just in case it's already gone +# original_message = None +# +# # If we should also delete the actual message, do that +# if delete_message and original_message is not None: +# try: +# await self.client.delete_message(original_message) +# except discord.Forbidden: +# self.logger.warning( +# "Did not have permission to remove message, aborting deletion (server={}, channel={})".format( +# channel.server.id, channel.id)) +# return +# +# # Remove it from the DB +# await db.delete_message(conn, message_id) +# +# # Then log deletion to logging channel +# await self.channel_logger.log_message_deleted(conn, +# server_id=channel.server.id, +# channel_name=channel.name, +# member_name=db_message.name, +# member_hid=db_message.hid, +# member_avatar_url=db_message.avatar_url, +# system_name=db_message.system_name, +# system_hid=db_message.system_hid, +# message_text=original_message.content if original_message else "*(original not found)*", +# message_id=message_id, +# deleted_by_moderator=deleted_by_moderator) +# +# async def handle_reaction(self, conn, user_id: str, message_id: str, emoji: str): +# if emoji == "❌": +# await self.try_delete_message(conn, message_id, check_user_id=user_id, delete_message=True, deleted_by_moderator=False) +# +# async def handle_deletion(self, conn, message_id: str): +# # Don't delete the message, it's already gone at this point, just handle DB deletion and logging +# await self.try_delete_message(conn, message_id, check_user_id=None, delete_message=False, deleted_by_moderator=True) diff --git a/src/pluralkit/bot/utils.py b/src/pluralkit/bot/utils.py index 789682d8..932de42e 100644 --- a/src/pluralkit/bot/utils.py +++ b/src/pluralkit/bot/utils.py @@ -2,6 +2,7 @@ import logging import re import discord +from typing import Optional from pluralkit import db from pluralkit.system import System @@ -20,28 +21,28 @@ def bounds_check_member_name(new_name, system_tag): if len("{} {}".format(new_name, system_tag)) > 32: return "This name, combined with the system tag ({}), would exceed the maximum length of 32 characters. Please reduce the length of the tag, or use a shorter name.".format(system_tag) -async def parse_mention(client: discord.Client, mention: str) -> discord.User: +async def parse_mention(client: discord.Client, mention: str) -> Optional[discord.User]: # First try matching mention format match = re.fullmatch("<@!?(\\d+)>", mention) if match: try: - return await client.get_user_info(match.group(1)) + return await client.get_user_info(int(match.group(1))) except discord.NotFound: return None # Then try with just ID try: - return await client.get_user_info(str(int(mention))) + return await client.get_user_info(int(mention)) except (ValueError, discord.NotFound): return None -def parse_channel_mention(mention: str, server: discord.Server) -> discord.Channel: +def parse_channel_mention(mention: str, server: discord.Guild) -> Optional[discord.TextChannel]: match = re.fullmatch("<#(\\d+)>", mention) if match: return server.get_channel(match.group(1)) try: - return server.get_channel(str(int(mention))) + return server.get_channel(int(mention)) except ValueError: return None diff --git a/src/pluralkit/db.py b/src/pluralkit/db.py index 747c1320..5b4faf73 100644 --- a/src/pluralkit/db.py +++ b/src/pluralkit/db.py @@ -62,17 +62,17 @@ async def delete_member(conn, member_id: int): @db_wrap -async def link_account(conn, system_id: int, account_id: str): +async def link_account(conn, system_id: int, account_id: int): logger.debug("Linking account (account_id={}, system_id={})".format( account_id, system_id)) - await conn.execute("insert into accounts (uid, system) values ($1, $2)", int(account_id), system_id) + await conn.execute("insert into accounts (uid, system) values ($1, $2)", account_id, system_id) @db_wrap -async def unlink_account(conn, system_id: int, account_id: str): +async def unlink_account(conn, system_id: int, account_id: int): logger.debug("Unlinking account (account_id={}, system_id={})".format( account_id, system_id)) - await conn.execute("delete from accounts where uid = $1 and system = $2", int(account_id), system_id) + await conn.execute("delete from accounts where uid = $1 and system = $2", account_id, system_id) @db_wrap @@ -81,8 +81,8 @@ async def get_linked_accounts(conn, system_id: int) -> List[int]: @db_wrap -async def get_system_by_account(conn, account_id: str) -> System: - row = await conn.fetchrow("select systems.* from systems, accounts where accounts.uid = $1 and accounts.system = systems.id", int(account_id)) +async def get_system_by_account(conn, account_id: int) -> System: + row = await conn.fetchrow("select systems.* from systems, accounts where accounts.uid = $1 and accounts.system = systems.id", account_id) return System(**row) if row else None @db_wrap @@ -151,26 +151,26 @@ async def get_members_exceeding(conn, system_id: int, length: int) -> List[Membe @db_wrap -async def get_webhook(conn, channel_id: str) -> (str, str): - row = await conn.fetchrow("select webhook, token from webhooks where channel = $1", int(channel_id)) +async def get_webhook(conn, channel_id: int) -> (str, str): + row = await conn.fetchrow("select webhook, token from webhooks where channel = $1", channel_id) return (str(row["webhook"]), row["token"]) if row else None @db_wrap -async def add_webhook(conn, channel_id: str, webhook_id: str, webhook_token: str): +async def add_webhook(conn, channel_id: int, webhook_id: int, webhook_token: str): logger.debug("Adding new webhook (channel={}, webhook={}, token={})".format( channel_id, webhook_id, webhook_token)) - await conn.execute("insert into webhooks (channel, webhook, token) values ($1, $2, $3)", int(channel_id), int(webhook_id), webhook_token) + await conn.execute("insert into webhooks (channel, webhook, token) values ($1, $2, $3)", channel_id, webhook_id, webhook_token) @db_wrap -async def delete_webhook(conn, channel_id: str): - await conn.execute("delete from webhooks where channel = $1", int(channel_id)) +async def delete_webhook(conn, channel_id: int): + await conn.execute("delete from webhooks where channel = $1", channel_id) @db_wrap -async def add_message(conn, message_id: str, channel_id: str, member_id: int, sender_id: str): +async def add_message(conn, message_id: int, channel_id: int, member_id: int, sender_id: int): logger.debug("Adding new message (id={}, channel={}, member={}, sender={})".format( message_id, channel_id, member_id, sender_id)) - await conn.execute("insert into messages (mid, channel, member, sender) values ($1, $2, $3, $4)", int(message_id), int(channel_id), member_id, int(sender_id)) + await conn.execute("insert into messages (mid, channel, member, sender) values ($1, $2, $3, $4)", message_id, channel_id, member_id, sender_id) class ProxyMember(namedtuple("ProxyMember", ["id", "hid", "prefix", "suffix", "color", "name", "avatar_url", "tag", "system_name", "system_hid"])): id: int @@ -185,7 +185,7 @@ class ProxyMember(namedtuple("ProxyMember", ["id", "hid", "prefix", "suffix", "c system_hid: str @db_wrap -async def get_members_by_account(conn, account_id: str) -> List[ProxyMember]: +async def get_members_by_account(conn, account_id: int) -> List[ProxyMember]: # Returns a "chimera" object rows = await conn.fetch("""select members.id, members.hid, members.prefix, members.suffix, members.color, members.name, members.avatar_url, @@ -195,7 +195,7 @@ async def get_members_by_account(conn, account_id: str) -> List[ProxyMember]: where accounts.uid = $1 and systems.id = accounts.system - and members.system = systems.id""", int(account_id)) + and members.system = systems.id""", account_id) return [ProxyMember(**row) for row in rows] class MessageInfo(namedtuple("MemberInfo", ["mid", "channel", "member", "sender", "name", "hid", "avatar_url", "system_name", "system_hid"])): @@ -220,7 +220,7 @@ class MessageInfo(namedtuple("MemberInfo", ["mid", "channel", "member", "sender" } @db_wrap -async def get_message_by_sender_and_id(conn, message_id: str, sender_id: str) -> MessageInfo: +async def get_message_by_sender_and_id(conn, message_id: int, sender_id: int) -> MessageInfo: row = await conn.fetchrow("""select messages.*, members.name, members.hid, members.avatar_url, @@ -231,12 +231,12 @@ async def get_message_by_sender_and_id(conn, message_id: str, sender_id: str) -> messages.member = members.id and members.system = systems.id and mid = $1 - and sender = $2""", int(message_id), int(sender_id)) + and sender = $2""", message_id, sender_id) return MessageInfo(**row) if row else None @db_wrap -async def get_message(conn, message_id: str) -> MessageInfo: +async def get_message(conn, message_id: int) -> MessageInfo: row = await conn.fetchrow("""select messages.*, members.name, members.hid, members.avatar_url, @@ -246,14 +246,14 @@ async def get_message(conn, message_id: str) -> MessageInfo: where messages.member = members.id and members.system = systems.id - and mid = $1""", int(message_id)) + and mid = $1""", message_id) return MessageInfo(**row) if row else None @db_wrap -async def delete_message(conn, message_id: str): +async def delete_message(conn, message_id: int): logger.debug("Deleting message (id={})".format(message_id)) - await conn.execute("delete from messages where mid = $1", int(message_id)) + await conn.execute("delete from messages where mid = $1", message_id) @db_wrap async def get_member_message_count(conn, member_id: int) -> int: @@ -290,14 +290,14 @@ async def add_switch_member(conn, switch_id: int, member_id: int): await conn.execute("insert into switch_members (switch, member) values ($1, $2)", switch_id, member_id) @db_wrap -async def get_server_info(conn, server_id: str): - return await conn.fetchrow("select * from servers where id = $1", int(server_id)) +async def get_server_info(conn, server_id: int): + return await conn.fetchrow("select * from servers where id = $1", server_id) @db_wrap -async def update_server(conn, server_id: str, logging_channel_id: str): - logging_channel_id = int(logging_channel_id) if logging_channel_id else None +async def update_server(conn, server_id: int, logging_channel_id: int): + logging_channel_id = logging_channel_id if logging_channel_id else None logger.debug("Updating server settings (id={}, log_channel={})".format(server_id, logging_channel_id)) - await conn.execute("insert into servers (id, log_channel) values ($1, $2) on conflict (id) do update set log_channel = $2", int(server_id), logging_channel_id) + await conn.execute("insert into servers (id, log_channel) values ($1, $2) on conflict (id) do update set log_channel = $2", server_id, logging_channel_id) @db_wrap async def member_count(conn) -> int: From ea62ede21b39cf372942bf407f507be6602ddc6e Mon Sep 17 00:00:00 2001 From: Ske Date: Sat, 27 Oct 2018 22:12:10 +0200 Subject: [PATCH 3/9] Basic command functionality --- src/pluralkit/bot/commands/__init__.py | 24 +++++++++++++----------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/src/pluralkit/bot/commands/__init__.py b/src/pluralkit/bot/commands/__init__.py index 6f346b28..b5940189 100644 --- a/src/pluralkit/bot/commands/__init__.py +++ b/src/pluralkit/bot/commands/__init__.py @@ -1,3 +1,4 @@ +import asyncio import discord import logging import re @@ -69,7 +70,7 @@ class CommandContext: def has_next(self) -> bool: return bool(self.args) - def pop_str(self, error: CommandError = None) -> str: + def pop_str(self, error: CommandError = None) -> Optional[str]: if not self.args: if error: raise error @@ -105,26 +106,27 @@ class CommandContext: return self.args async def reply(self, content=None, embed=None): - return await self.client.send_message(self.message.channel, content=content, embed=embed) + return await self.message.channel.send(content=content, embed=embed) async def confirm_react(self, user: Union[discord.Member, discord.User], message: str): message = await self.reply(message) + await message.add_reaction("\u2705") # Checkmark + await message.add_reaction("\u274c") # Red X - await self.client.add_reaction(message, "✅") - await self.client.add_reaction(message, "❌") - - reaction = await self.client.wait_for_reaction(emoji=["✅", "❌"], user=user, timeout=60.0*5) - if not reaction: + try: + reaction, _ = await self.client.wait_for("reaction_add", check=lambda r, u: u.id == user.id and r.emoji in ["\u2705", "\u274c"], timeout=60.0*5) + return reaction.emoji == "\u2705" + except asyncio.TimeoutError: raise CommandError("Timed out - try again.") - return reaction.reaction.emoji == "✅" async def confirm_text(self, user: discord.Member, channel: discord.TextChannel, confirm_text: str, message: str): await self.reply(message) - message = await self.client.wait_for_message(channel=channel, author=user, timeout=60.0*5) - if not message: + try: + message = await self.client.wait_for("message", check=lambda m: m.channel.id == channel.id and m.author.id == user.id, timeout=60.0*5) + return message.content.lower() == confirm_text.lower() + except asyncio.TimeoutError: raise CommandError("Timed out - try again.") - return message.content == confirm_text import pluralkit.bot.commands.import_commands From 58d8927380f540bc93d3a0e684f927a95c955f8b Mon Sep 17 00:00:00 2001 From: Ske Date: Sat, 27 Oct 2018 23:30:12 +0200 Subject: [PATCH 4/9] Message and error logging, various bugfixes --- src/pluralkit/bot/__init__.py | 62 +++++++++++++- src/pluralkit/bot/channel_logger.py | 32 ++++---- .../bot/commands/message_commands.py | 6 +- src/pluralkit/bot/commands/mod_commands.py | 7 +- src/pluralkit/bot/embeds.py | 8 +- src/pluralkit/bot/proxy.py | 81 +++++++++++++++++-- src/pluralkit/bot/utils.py | 2 +- 7 files changed, 161 insertions(+), 37 deletions(-) diff --git a/src/pluralkit/bot/__init__.py b/src/pluralkit/bot/__init__.py index 1d2708bc..ee1f60cd 100644 --- a/src/pluralkit/bot/__init__.py +++ b/src/pluralkit/bot/__init__.py @@ -7,12 +7,14 @@ import os import logging import discord +import traceback from pluralkit import db -from pluralkit.bot import commands, proxy +from pluralkit.bot import commands, proxy, channel_logger, embeds logging.basicConfig(level=logging.INFO, format="[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s") + def connect_to_database() -> asyncpg.pool.Pool: username = os.environ["DATABASE_USER"] password = os.environ["DATABASE_PASS"] @@ -21,7 +23,9 @@ def connect_to_database() -> asyncpg.pool.Pool: port = os.environ["DATABASE_PORT"] if username is None or password is None or name is None or host is None or port is None: - print("Database credentials not specified. Please pass valid PostgreSQL database credentials in the DATABASE_[USER|PASS|NAME|HOST|PORT] environment variable.", file=sys.stderr) + print( + "Database credentials not specified. Please pass valid PostgreSQL database credentials in the DATABASE_[USER|PASS|NAME|HOST|PORT] environment variable.", + file=sys.stderr) sys.exit(1) try: @@ -38,11 +42,20 @@ def connect_to_database() -> asyncpg.pool.Pool: port=port )) + def run(): pool = connect_to_database() + async def create_tables(): + async with pool.acquire() as conn: + await db.create_tables(conn) + + asyncio.get_event_loop().run_until_complete(create_tables()) + client = discord.Client() + logger = channel_logger.ChannelLogger(client) + @client.event async def on_ready(): print("PluralKit started.") @@ -64,8 +77,51 @@ def run(): return # Second pass: do proxy matching - await proxy.try_proxy_message(message, conn) + await proxy.try_proxy_message(conn, message, logger) + @client.event + async def on_raw_message_delete(payload: discord.RawMessageDeleteEvent): + async with pool.acquire() as conn: + await proxy.handle_deleted_message(conn, client, payload.message_id, None, logger) + + @client.event + async def on_raw_bulk_message_delete(payload: discord.RawBulkMessageDeleteEvent): + async with pool.acquire() as conn: + for message_id in payload.message_ids: + await proxy.handle_deleted_message(conn, client, message_id, None, logger) + + @client.event + async def on_raw_reaction_add(payload: discord.RawReactionActionEvent): + if payload.emoji.name == "\u274c": # Red X + async with pool.acquire() as conn: + await proxy.try_delete_by_reaction(conn, client, payload.message_id, payload.user_id, logger) + + @client.event + async def on_error(event_name, *args, **kwargs): + log_channel_id = os.environ["LOG_CHANNEL"] + if not log_channel_id: + return + + log_channel = client.get_channel(int(log_channel_id)) + + # If this is a message event, we can attach additional information in an event + # ie. username, channel, content, etc + if args and isinstance(args[0], discord.Message): + message: discord.Message = args[0] + embed = embeds.exception_log( + message.content, + message.author.name, + message.author.discriminator, + message.author.id, + message.guild.id if message.guild else None, + message.channel.id + ) + else: + # If not, just post the string itself + embed = None + + traceback_str = "```python\n{}```".format(traceback.format_exc()) + await log_channel.send(content=traceback_str, embed=embed) bot_token = os.environ["TOKEN"] if not bot_token: diff --git a/src/pluralkit/bot/channel_logger.py b/src/pluralkit/bot/channel_logger.py index 7285cced..acbb2415 100644 --- a/src/pluralkit/bot/channel_logger.py +++ b/src/pluralkit/bot/channel_logger.py @@ -19,7 +19,7 @@ class ChannelLogger: self.logger = logging.getLogger("pluralkit.bot.channel_logger") self.client = client - async def get_log_channel(self, conn, server_id: str): + async def get_log_channel(self, conn, server_id: int): server_info = await db.get_server_info(conn, server_id) if not server_info: @@ -30,21 +30,21 @@ class ChannelLogger: if not log_channel: return None - return self.client.get_channel(str(log_channel)) + return self.client.get_channel(log_channel) - async def send_to_log_channel(self, log_channel: discord.Channel, embed: discord.Embed, text: str = None): + async def send_to_log_channel(self, log_channel: discord.TextChannel, embed: discord.Embed, text: str = None): try: - await self.client.send_message(log_channel, embed=embed, content=text) + await log_channel.send(content=text, embed=embed) except discord.Forbidden: # TODO: spew big error self.logger.warning( "Did not have permission to send message to logging channel (server={}, channel={})".format( - log_channel.server.id, log_channel.id)) + log_channel.guild.id, log_channel.id)) async def log_message_proxied(self, conn, - server_id: str, + server_id: int, channel_name: str, - channel_id: str, + channel_id: int, sender_name: str, sender_disc: int, sender_id: int, @@ -56,11 +56,13 @@ class ChannelLogger: message_text: str, message_image: str, message_timestamp: datetime, - message_id: str): + message_id: int): log_channel = await self.get_log_channel(conn, server_id) if not log_channel: return + message_link = "https://discordapp.com/channels/{}/{}/{}".format(server_id, channel_id, message_id) + embed = discord.Embed() embed.colour = discord.Colour.blue() embed.description = message_text @@ -75,11 +77,10 @@ class ChannelLogger: if message_image: embed.set_thumbnail(url=message_image) - message_link = "https://discordapp.com/channels/{}/{}/{}".format(server_id, channel_id, message_id) await self.send_to_log_channel(log_channel, embed, message_link) async def log_message_deleted(self, conn, - server_id: str, + server_id: int, channel_name: str, member_name: str, member_hid: str, @@ -87,22 +88,17 @@ class ChannelLogger: system_name: str, system_hid: str, message_text: str, - message_id: str, - deleted_by_moderator: bool): + message_id: int): log_channel = await self.get_log_channel(conn, server_id) if not log_channel: return embed = discord.Embed() embed.colour = discord.Colour.dark_red() - embed.description = message_text + embed.description = message_text or "*(unknown, message deleted by moderator)*" embed.timestamp = datetime.utcnow() embed_set_author_name(embed, channel_name, member_name, system_name, member_avatar_url) - embed.set_footer( - text="System ID: {} | Member ID: {} | Message ID: {} | Deleted by moderator? {}".format(system_hid, - member_hid, - message_id, - "Yes" if deleted_by_moderator else "No")) + embed.set_footer(text="System ID: {} | Member ID: {} | Message ID: {}".format(system_hid, member_hid, message_id)) await self.send_to_log_channel(log_channel, embed) diff --git a/src/pluralkit/bot/commands/message_commands.py b/src/pluralkit/bot/commands/message_commands.py index 68909ef6..7d474d8e 100644 --- a/src/pluralkit/bot/commands/message_commands.py +++ b/src/pluralkit/bot/commands/message_commands.py @@ -8,7 +8,7 @@ async def get_message_contents(client: discord.Client, channel_id: int, message_ channel = client.get_channel(str(channel_id)) if channel: try: - original_message = await client.get_message(channel, str(message_id)) + original_message = await client.get_channel(channel).get_message(message_id) return original_message.content or None except (discord.errors.Forbidden, discord.errors.NotFound): pass @@ -24,13 +24,13 @@ async def message_info(ctx: CommandContext): return CommandError("You must pass a valid number as a message ID.", help=help.message_lookup) # Find the message in the DB - message = await db.get_message(ctx.conn, str(mid)) + message = await db.get_message(ctx.conn, mid) if not message: raise CommandError("Message with ID '{}' not found.".format(mid)) # Get the original sender of the messages try: - original_sender = await ctx.client.get_user_info(str(message.sender)) + original_sender = await ctx.client.get_user_info(message.sender) except discord.NotFound: # Account was since deleted - rare but we're handling it anyway original_sender = None diff --git a/src/pluralkit/bot/commands/mod_commands.py b/src/pluralkit/bot/commands/mod_commands.py index 7804113e..8d532825 100644 --- a/src/pluralkit/bot/commands/mod_commands.py +++ b/src/pluralkit/bot/commands/mod_commands.py @@ -4,10 +4,13 @@ logger = logging.getLogger("pluralkit.commands") async def set_log(ctx: CommandContext): - if not ctx.message.author.server_permissions.administrator: + if not ctx.message.author.guild_permissions.administrator: return CommandError("You must be a server administrator to use this command.") - server = ctx.message.server + server = ctx.message.guild + if not server: + return CommandError("This command can not be run in a DM.") + if not ctx.has_next(): channel_id = None else: diff --git a/src/pluralkit/bot/embeds.py b/src/pluralkit/bot/embeds.py index b134bb1c..ec5bd038 100644 --- a/src/pluralkit/bot/embeds.py +++ b/src/pluralkit/bot/embeds.py @@ -38,13 +38,13 @@ def status(text: str) -> discord.Embed: return embed -def exception_log(message_content, author_name, author_discriminator, server_id, channel_id) -> discord.Embed: +def exception_log(message_content, author_name, author_discriminator, author_id, server_id, channel_id) -> discord.Embed: embed = discord.Embed() embed.colour = discord.Colour.dark_red() embed.title = message_content - embed.set_footer(text="Sender: {}#{} | Server: {} | Channel: {}".format( - author_name, author_discriminator, + embed.set_footer(text="Sender: {}#{} ({}) | Server: {} | Channel: {}".format( + author_name, author_discriminator, author_id, server_id if server_id else "(DMs)", channel_id )) @@ -72,7 +72,7 @@ async def system_card(conn, client: discord.Client, system: System) -> discord.E account_names = [] for account_id in await system.get_linked_account_ids(conn): - account = await client.get_user_info(str(account_id)) + account = await client.get_user_info(account_id) account_names.append("{}#{}".format(account.name, account.discriminator)) card.add_field(name="Linked accounts", value="\n".join(account_names)) diff --git a/src/pluralkit/bot/proxy.py b/src/pluralkit/bot/proxy.py index a6920531..22d0f45c 100644 --- a/src/pluralkit/bot/proxy.py +++ b/src/pluralkit/bot/proxy.py @@ -3,10 +3,11 @@ from io import BytesIO import discord import logging import re -from typing import List +from typing import List, Optional from pluralkit import db -from pluralkit.bot import utils +from pluralkit.bot import utils, channel_logger +from pluralkit.bot.channel_logger import ChannelLogger logger = logging.getLogger("pluralkit.bot.proxy") @@ -99,7 +100,7 @@ async def make_attachment_file(message: discord.Message): async def do_proxy_message(conn, original_message: discord.Message, proxy_member: db.ProxyMember, - inner_text: str): + inner_text: str, logger: ChannelLogger): # Send the message through the webhook webhook = await get_or_create_webhook_for_channel(conn, original_message.channel) sent_message = await webhook.send( @@ -114,10 +115,30 @@ async def do_proxy_message(conn, original_message: discord.Message, proxy_member await db.add_message(conn, sent_message.id, original_message.channel.id, proxy_member.id, original_message.author.id) - # TODO: log message in log channel + await logger.log_message_proxied( + conn, + original_message.channel.guild.id, + original_message.channel.name, + original_message.channel.id, + original_message.author.name, + original_message.author.discriminator, + original_message.author.id, + proxy_member.name, + proxy_member.hid, + proxy_member.avatar_url, + proxy_member.system_name, + proxy_member.system_hid, + inner_text, + sent_message.attachments[0].url if sent_message.attachments else None, + sent_message.created_at, + sent_message.id + ) + + # And finally, gotta delete the original. + await original_message.delete() -async def try_proxy_message(message: discord.Message, conn) -> bool: +async def try_proxy_message(conn, message: discord.Message, logger: ChannelLogger) -> bool: # Don't bother proxying in DMs with the bot if isinstance(message.channel, discord.abc.PrivateChannel): return False @@ -140,9 +161,57 @@ async def try_proxy_message(message: discord.Message, conn) -> bool: # So, we now have enough information to successfully proxy a message async with conn.transaction(): - await do_proxy_message(conn, message, member, inner_text) + await do_proxy_message(conn, message, member, inner_text, logger) return True + +async def handle_deleted_message(conn, client: discord.Client, message_id: int, + message_content: Optional[str], logger: channel_logger.ChannelLogger) -> bool: + msg = await db.get_message(conn, message_id) + if not msg: + return False + + channel = client.get_channel(msg.channel) + if not channel: + # Weird edge case, but channel *could* be deleted at this point (can't think of any scenarios it would be tho) + return False + + await db.delete_message(conn, message_id) + await logger.log_message_deleted( + conn, + channel.guild.id, + channel.name, + msg.name, + msg.hid, + msg.avatar_url, + msg.system_name, + msg.system_hid, + message_content, + message_id + ) + return True + + +async def try_delete_by_reaction(conn, client: discord.Client, message_id: int, reaction_user: int, + logger: channel_logger.ChannelLogger) -> bool: + # Find the message by the given message id or reaction user + msg = await db.get_message_by_sender_and_id(conn, message_id, reaction_user) + if not msg: + # Either the wrong user reacted or the message isn't a proxy message + # In either case - not our problem + return False + + # Find the original message + original_message = await client.get_channel(msg.channel).get_message(message_id) + if not original_message: + # Message got deleted, possibly race condition, eh + return False + + # Then delete the original message + await original_message.delete() + + await handle_deleted_message(conn, client, message_id, original_message.content, logger) + # # TODO: possibly move this to bot __init__ so commands can access it too # class WebhookPermissionError(Exception): # pass diff --git a/src/pluralkit/bot/utils.py b/src/pluralkit/bot/utils.py index 932de42e..8060168d 100644 --- a/src/pluralkit/bot/utils.py +++ b/src/pluralkit/bot/utils.py @@ -39,7 +39,7 @@ async def parse_mention(client: discord.Client, mention: str) -> Optional[discor def parse_channel_mention(mention: str, server: discord.Guild) -> Optional[discord.TextChannel]: match = re.fullmatch("<#(\\d+)>", mention) if match: - return server.get_channel(match.group(1)) + return server.get_channel(int(match.group(1))) try: return server.get_channel(int(mention)) From a4449ceaae4ca4cfb3adaa564a8913939df6e365 Mon Sep 17 00:00:00 2001 From: Ske Date: Sat, 27 Oct 2018 23:31:45 +0200 Subject: [PATCH 5/9] Remove commented-out old code --- src/pluralkit/bot/__init__.py | 113 +---------------- src/pluralkit/bot/proxy.py | 231 +--------------------------------- 2 files changed, 2 insertions(+), 342 deletions(-) diff --git a/src/pluralkit/bot/__init__.py b/src/pluralkit/bot/__init__.py index ee1f60cd..b685b21d 100644 --- a/src/pluralkit/bot/__init__.py +++ b/src/pluralkit/bot/__init__.py @@ -129,115 +129,4 @@ def run(): file=sys.stderr) sys.exit(1) - client.run(bot_token) - -# logging.getLogger("pluralkit").setLevel(logging.DEBUG) - -# class PluralKitBot: -# def __init__(self, token): -# self.token = token -# self.logger = logging.getLogger("pluralkit.bot") -# -# self.client = discord.Client() -# self.client.event(self.on_error) -# self.client.event(self.on_ready) -# self.client.event(self.on_message) -# self.client.event(self.on_socket_raw_receive) -# -# self.channel_logger = channel_logger.ChannelLogger(self.client) -# -# self.proxy = proxy.Proxy(self.client, token, self.channel_logger) -# -# async def on_error(self, evt, *args, **kwargs): -# self.logger.exception("Error while handling event {} with arguments {}:".format(evt, args)) -# -# async def on_ready(self): -# self.logger.info("Connected to Discord.") -# self.logger.info("- Account: {}#{}".format(self.client.user.name, self.client.user.discriminator)) -# self.logger.info("- User ID: {}".format(self.client.user.id)) -# self.logger.info("- {} servers".format(len(self.client.servers))) -# -# async def on_message(self, message): -# # Ignore bot messages -# if message.author.bot: -# return -# -# try: -# if await self.handle_command_dispatch(message): -# return -# -# if await self.handle_proxy_dispatch(message): -# return -# except Exception: -# await self.log_error_in_channel(message) -# -# async def on_socket_raw_receive(self, msg): -# # Since on_reaction_add is buggy (only works for messages the bot's already cached, ie. no old messages) -# # we parse socket data manually for the reaction add event -# if isinstance(msg, str): -# try: -# msg_data = json.loads(msg) -# if msg_data.get("t") == "MESSAGE_REACTION_ADD": -# evt_data = msg_data.get("d") -# if evt_data: -# user_id = evt_data["user_id"] -# message_id = evt_data["message_id"] -# emoji = evt_data["emoji"]["name"] -# -# async with self.pool.acquire() as conn: -# await self.proxy.handle_reaction(conn, user_id, message_id, emoji) -# elif msg_data.get("t") == "MESSAGE_DELETE": -# evt_data = msg_data.get("d") -# if evt_data: -# message_id = evt_data["id"] -# async with self.pool.acquire() as conn: -# await self.proxy.handle_deletion(conn, message_id) -# except ValueError: -# pass -# -# async def handle_command_dispatch(self, message): -# async with self.pool.acquire() as conn: -# result = await commands.command_dispatch(self.client, message, conn) -# return result -# -# async def handle_proxy_dispatch(self, message): -# # Try doing proxy parsing -# async with self.pool.acquire() as conn: -# return await self.proxy.try_proxy_message(conn, message) -# -# async def log_error_in_channel(self, message): -# channel_id = os.environ["LOG_CHANNEL"] -# if not channel_id: -# return -# -# channel = self.client.get_channel(channel_id) -# embed = embeds.exception_log( -# message.content, -# message.author.name, -# message.author.discriminator, -# message.server.id if message.server else None, -# message.channel.id -# ) -# -# await self.client.send_message(channel, "```python\n{}```".format(traceback.format_exc()), embed=embed) -# -# async def run(self): -# try: -# self.logger.info("Connecting to database...") -# self.pool = await db.connect( -# os.environ["DATABASE_USER"], -# os.environ["DATABASE_PASS"], -# os.environ["DATABASE_NAME"], -# os.environ["DATABASE_HOST"], -# int(os.environ["DATABASE_PORT"]) -# ) -# -# self.logger.info("Attempting to create tables...") -# async with self.pool.acquire() as conn: -# await db.create_tables(conn) -# -# self.logger.info("Connecting to Discord...") -# await self.client.start(self.token) -# finally: -# self.logger.info("Logging out from Discord...") -# await self.client.logout() + client.run(bot_token) \ No newline at end of file diff --git a/src/pluralkit/bot/proxy.py b/src/pluralkit/bot/proxy.py index 22d0f45c..8997a948 100644 --- a/src/pluralkit/bot/proxy.py +++ b/src/pluralkit/bot/proxy.py @@ -210,233 +210,4 @@ async def try_delete_by_reaction(conn, client: discord.Client, message_id: int, # Then delete the original message await original_message.delete() - await handle_deleted_message(conn, client, message_id, original_message.content, logger) - -# # TODO: possibly move this to bot __init__ so commands can access it too -# class WebhookPermissionError(Exception): -# pass -# -# -# class DeletionPermissionError(Exception): -# pass -# -# -# class Proxy: -# def __init__(self, client: discord.Client, token: str, logger: channel_logger.ChannelLogger): -# self.logger = logging.getLogger("pluralkit.bot.proxy") -# self.session = aiohttp.ClientSession() -# self.client = client -# self.token = token -# self.channel_logger = logger -# -# async def save_channel_webhook(self, conn, channel: discord.Channel, id: str, token: str) -> (str, str): -# await db.add_webhook(conn, channel.id, id, token) -# return id, token -# -# async def create_and_add_channel_webhook(self, conn, channel: discord.Channel) -> (str, str): -# # This method is only called if there's no webhook found in the DB (and hopefully within a transaction) -# # No need to worry about error handling if there's a DB conflict (which will throw an exception because DB constraints) -# req_headers = {"Authorization": "Bot {}".format(self.token)} -# -# # First, check if there's already a webhook belonging to the bot -# async with self.session.get("https://discordapp.com/api/v6/channels/{}/webhooks".format(channel.id), -# headers=req_headers) as resp: -# if resp.status == 200: -# webhooks = await resp.json() -# for webhook in webhooks: -# if webhook["user"]["id"] == self.client.user.id: -# # This webhook belongs to us, we can use that, return it and save it -# return await self.save_channel_webhook(conn, channel, webhook["id"], webhook["token"]) -# elif resp.status == 403: -# self.logger.warning( -# "Did not have permission to fetch webhook list (server={}, channel={})".format(channel.server.id, -# channel.id)) -# raise WebhookPermissionError() -# else: -# raise discord.HTTPException(resp, await resp.text()) -# -# # Then, try submitting a new one -# req_data = {"name": "PluralKit Proxy Webhook"} -# async with self.session.post("https://discordapp.com/api/v6/channels/{}/webhooks".format(channel.id), -# json=req_data, headers=req_headers) as resp: -# if resp.status == 200: -# webhook = await resp.json() -# return await self.save_channel_webhook(conn, channel, webhook["id"], webhook["token"]) -# elif resp.status == 403: -# self.logger.warning( -# "Did not have permission to create webhook (server={}, channel={})".format(channel.server.id, -# channel.id)) -# raise WebhookPermissionError() -# else: -# raise discord.HTTPException(resp, await resp.text()) -# -# # Should not be reached without an exception being thrown -# -# async def get_webhook_for_channel(self, conn, channel: discord.Channel): -# async with conn.transaction(): -# hook_match = await db.get_webhook(conn, channel.id) -# if not hook_match: -# # We don't have a webhook, create/add one -# return await self.create_and_add_channel_webhook(conn, channel) -# else: -# return hook_match -# -# async def do_proxy_message(self, conn, member: db.ProxyMember, original_message: discord.Message, text: str, -# attachment_url: str, has_already_retried=False): -# hook_id, hook_token = await self.get_webhook_for_channel(conn, original_message.channel) -# -# form_data = aiohttp.FormData() -# form_data.add_field("username", "{} {}".format(member.name, member.tag or "").strip()) -# -# if text: -# form_data.add_field("content", text) -# -# if attachment_url: -# attachment_resp = await self.session.get(attachment_url) -# form_data.add_field("file", attachment_resp.content, content_type=attachment_resp.content_type, -# filename=attachment_resp.url.name) -# -# if member.avatar_url: -# form_data.add_field("avatar_url", member.avatar_url) -# -# async with self.session.post( -# "https://discordapp.com/api/v6/webhooks/{}/{}?wait=true".format(hook_id, hook_token), -# data=form_data) as resp: -# if resp.status == 200: -# message = await resp.json() -# -# await db.add_message(conn, message["id"], message["channel_id"], member.id, original_message.author.id) -# -# try: -# await self.client.delete_message(original_message) -# except discord.Forbidden: -# self.logger.warning( -# "Did not have permission to delete original message (server={}, channel={})".format( -# original_message.server.id, original_message.channel.id)) -# raise DeletionPermissionError() -# except discord.NotFound: -# self.logger.warning("Tried to delete message when proxying, but message was already gone (server={}, channel={})".format(original_message.server.id, original_message.channel.id)) -# -# message_image = None -# if message["attachments"]: -# first_attachment = message["attachments"][0] -# if "width" in first_attachment and "height" in first_attachment: -# # Only log attachments that are actually images -# message_image = first_attachment["url"] -# -# await self.channel_logger.log_message_proxied(conn, -# server_id=original_message.server.id, -# channel_name=original_message.channel.name, -# channel_id=original_message.channel.id, -# sender_name=original_message.author.name, -# sender_disc=original_message.author.discriminator, -# sender_id=original_message.author.id, -# member_name=member.name, -# member_hid=member.hid, -# member_avatar_url=member.avatar_url, -# system_name=member.system_name, -# system_hid=member.system_hid, -# message_text=text, -# message_image=message_image, -# message_timestamp=ciso8601.parse_datetime( -# message["timestamp"]), -# message_id=message["id"]) -# elif resp.status == 404 and not has_already_retried: -# # Webhook doesn't exist. Delete it from the DB, create, and add a new one -# self.logger.warning("Webhook registered in DB doesn't exist, deleting hook from DB, re-adding, and trying again (channel={}, hook={})".format(original_message.channel.id, hook_id)) -# await db.delete_webhook(conn, original_message.channel.id) -# await self.create_and_add_channel_webhook(conn, original_message.channel) -# -# # Then try again all over, making sure to not retry again and go in a loop should it continually fail -# return await self.do_proxy_message(conn, member, original_message, text, attachment_url, has_already_retried=True) -# else: -# raise discord.HTTPException(resp, await resp.text()) -# -# async def try_proxy_message(self, conn, message: discord.Message): -# # Can't proxy in DMs, webhook creation will explode -# if message.channel.is_private: -# return False -# -# # Big fat query to find every member associated with this account -# # Returned member object has a few more keys (system tag, for example) -# members = await db.get_members_by_account(conn, account_id=message.author.id) -# -# match = match_proxy_tags(members, message.content) -# if not match: -# return False -# -# member, text = match -# attachment_url = get_message_attachment_url(message) -# -# # Can't proxy a message with no text AND no attachment -# if not text and not attachment_url: -# self.logger.debug("Skipping message because of no text and no attachment") -# return False -# -# # Remember to sanitize the text (remove @everyones and such) -# text = utils.sanitize(text) -# -# try: -# async with conn.transaction(): -# await self.do_proxy_message(conn, member, message, text=text, attachment_url=attachment_url) -# except WebhookPermissionError: -# embed = embeds.error("PluralKit does not have permission to manage webhooks for this channel. Contact your local server administrator to fix this.") -# await self.client.send_message(message.channel, embed=embed) -# except DeletionPermissionError: -# embed = embeds.error("PluralKit does not have permission to delete messages in this channel. Contact your local server administrator to fix this.") -# await self.client.send_message(message.channel, embed=embed) -# -# return True -# -# async def try_delete_message(self, conn, message_id: str, check_user_id: Optional[str], delete_message: bool, deleted_by_moderator: bool): -# async with conn.transaction(): -# # Find the message in the DB, and make sure it's sent by the user (if we need to check) -# if check_user_id: -# db_message = await db.get_message_by_sender_and_id(conn, message_id=message_id, sender_id=check_user_id) -# else: -# db_message = await db.get_message(conn, message_id=message_id) -# -# if db_message: -# self.logger.debug("Deleting message {}".format(message_id)) -# channel = self.client.get_channel(str(db_message.channel)) -# -# # Fetch the original message from the server -# try: -# original_message = await self.client.get_message(channel, message_id) -# except discord.NotFound: -# # Just in case it's already gone -# original_message = None -# -# # If we should also delete the actual message, do that -# if delete_message and original_message is not None: -# try: -# await self.client.delete_message(original_message) -# except discord.Forbidden: -# self.logger.warning( -# "Did not have permission to remove message, aborting deletion (server={}, channel={})".format( -# channel.server.id, channel.id)) -# return -# -# # Remove it from the DB -# await db.delete_message(conn, message_id) -# -# # Then log deletion to logging channel -# await self.channel_logger.log_message_deleted(conn, -# server_id=channel.server.id, -# channel_name=channel.name, -# member_name=db_message.name, -# member_hid=db_message.hid, -# member_avatar_url=db_message.avatar_url, -# system_name=db_message.system_name, -# system_hid=db_message.system_hid, -# message_text=original_message.content if original_message else "*(original not found)*", -# message_id=message_id, -# deleted_by_moderator=deleted_by_moderator) -# -# async def handle_reaction(self, conn, user_id: str, message_id: str, emoji: str): -# if emoji == "❌": -# await self.try_delete_message(conn, message_id, check_user_id=user_id, delete_message=True, deleted_by_moderator=False) -# -# async def handle_deletion(self, conn, message_id: str): -# # Don't delete the message, it's already gone at this point, just handle DB deletion and logging -# await self.try_delete_message(conn, message_id, check_user_id=None, delete_message=False, deleted_by_moderator=True) + await handle_deleted_message(conn, client, message_id, original_message.content, logger) \ No newline at end of file From 899183f26d2904e7970acbb43c2ef3e4f4ca8f80 Mon Sep 17 00:00:00 2001 From: Ske Date: Sat, 27 Oct 2018 23:34:50 +0200 Subject: [PATCH 6/9] Fix special case where webhook is deleted by user --- src/pluralkit/bot/proxy.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/src/pluralkit/bot/proxy.py b/src/pluralkit/bot/proxy.py index 8997a948..a025f99d 100644 --- a/src/pluralkit/bot/proxy.py +++ b/src/pluralkit/bot/proxy.py @@ -103,13 +103,21 @@ async def do_proxy_message(conn, original_message: discord.Message, proxy_member inner_text: str, logger: ChannelLogger): # Send the message through the webhook webhook = await get_or_create_webhook_for_channel(conn, original_message.channel) - sent_message = await webhook.send( - content=inner_text, - username=proxy_member.name, - avatar_url=proxy_member.avatar_url, - file=await make_attachment_file(original_message), - wait=True - ) + + try: + sent_message = await webhook.send( + content=inner_text, + username=proxy_member.name, + avatar_url=proxy_member.avatar_url, + file=await make_attachment_file(original_message), + wait=True + ) + except discord.NotFound: + # The webhook we got from the DB doesn't actually exist + # If we delete it from the DB then call the function again, it'll re-create one for us + await db.delete_webhook(conn, original_message.channel.id) + await do_proxy_message(conn, original_message, proxy_member, inner_text, logger) + return # Save the proxied message in the database await db.add_message(conn, sent_message.id, original_message.channel.id, proxy_member.id, From 99657ae07dd5818da2c81bff68daa685a790afee Mon Sep 17 00:00:00 2001 From: Ske Date: Sat, 27 Oct 2018 23:47:28 +0200 Subject: [PATCH 7/9] Fix Tupperware importing --- src/pluralkit/bot/commands/import_commands.py | 43 ++++++++----------- 1 file changed, 19 insertions(+), 24 deletions(-) diff --git a/src/pluralkit/bot/commands/import_commands.py b/src/pluralkit/bot/commands/import_commands.py index e74805e7..6c6fa2d5 100644 --- a/src/pluralkit/bot/commands/import_commands.py +++ b/src/pluralkit/bot/commands/import_commands.py @@ -8,22 +8,16 @@ logger = logging.getLogger("pluralkit.commands") async def import_tupperware(ctx: CommandContext): - tupperware_ids = ["431544605209788416", "433916057053560832"] # Main bot instance and Multi-Pals-specific fork - tupperware_members = [ctx.message.server.get_member(bot_id) for bot_id in tupperware_ids if - ctx.message.server.get_member(bot_id)] + tupperware_member = ctx.message.guild.get_member(431544605209788416) - # Check if there's any Tupperware bot on the server - if not tupperware_members: + # Check if there's a Tupperware bot on the server + if not tupperware_member: return CommandError("This command only works in a server where the Tupperware bot is also present.") - # Make sure at least one of the bts have send/read permissions here - for bot_member in tupperware_members: - channel_permissions = ctx.message.channel.permissions_for(bot_member) - if channel_permissions.read_messages and channel_permissions.send_messages: - # If so, break out of the loop - break - else: - # If no bots have permission (ie. loop doesn't break), throw error + # Make sure at the bot has send/read permissions here + channel_permissions = ctx.message.channel.permissions_for(tupperware_member) + if not (channel_permissions.read_messages and channel_permissions.send_messages): + # If it doesn't, throw error return CommandError("This command only works in a channel where the Tupperware bot has read/send access.") await ctx.reply( @@ -31,29 +25,31 @@ async def import_tupperware(ctx: CommandContext): # Check to make sure the message is sent by Tupperware, and that the Tupperware response actually belongs to the correct user def ensure_account(tw_msg): - if tw_msg.author not in tupperware_members: + if tw_msg.channel.id != ctx.message.channel.id: + return False + + if tw_msg.author.id != tupperware_member.id: return False if not tw_msg.embeds: return False - if not tw_msg.embeds[0]["title"]: + if not tw_msg.embeds[0].title: return False - return tw_msg.embeds[0]["title"].startswith( + return tw_msg.embeds[0].title.startswith( "{}#{}".format(ctx.message.author.name, ctx.message.author.discriminator)) tupperware_page_embeds = [] - tw_msg: discord.Message = await ctx.client.wait_for_message(channel=ctx.message.channel, timeout=60.0 * 5, - check=ensure_account) + tw_msg: discord.Message = await ctx.client.wait_for("message", check=ensure_account, timeout=60.0 * 5) if not tw_msg: return CommandError("Tupperware import timed out.") - tupperware_page_embeds.append(tw_msg.embeds[0]) + tupperware_page_embeds.append(tw_msg.embeds[0].to_dict()) # Handle Tupperware pagination def match_pagination(): - pagination_match = re.search(r"\(page (\d+)/(\d+), \d+ total\)", tw_msg.embeds[0]["title"]) + pagination_match = re.search(r"\(page (\d+)/(\d+), \d+ total\)", tw_msg.embeds[0].title) if not pagination_match: return None return int(pagination_match.group(1)), int(pagination_match.group(2)) @@ -72,13 +68,12 @@ async def import_tupperware(ctx: CommandContext): new_page, total_pages = match_pagination() # Put the found page in the pages dict - pages_found[new_page] = dict(tw_msg.embeds[0]) + pages_found[new_page] = tw_msg.embeds[0].to_dict() # If this isn't the same page as last check, edit the status message if new_page != current_page: last_found_time = datetime.utcnow() - await ctx.client.edit_message(status_msg, - "Multi-page member list found. Please manually scroll through all the pages. Read {}/{} pages.".format( + await status_msg.edit(content="Multi-page member list found. Please manually scroll through all the pages. Read {}/{} pages.".format( len(pages_found), total_pages)) current_page = new_page @@ -94,7 +89,7 @@ async def import_tupperware(ctx: CommandContext): tupperware_page_embeds = list([embed for page, embed in sorted(pages_found.items(), key=lambda x: x[0])]) # Also edit the status message to indicate we're now importing, and it may take a while because there's probably a lot of members - await ctx.client.edit_message(status_msg, "All pages read. Now importing...") + await status_msg.edit(content="All pages read. Now importing...") logger.debug("Importing from Tupperware...") From 3fca071c1d3c79da2bbf6c905f27255621f21466 Mon Sep 17 00:00:00 2001 From: Ske Date: Thu, 8 Nov 2018 15:47:21 +0100 Subject: [PATCH 8/9] Count all-blank strings as empty for the purposes of proxy validation --- src/pluralkit/bot/proxy.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/pluralkit/bot/proxy.py b/src/pluralkit/bot/proxy.py index a025f99d..078f250f 100644 --- a/src/pluralkit/bot/proxy.py +++ b/src/pluralkit/bot/proxy.py @@ -164,7 +164,8 @@ async def try_proxy_message(conn, message: discord.Message, logger: ChannelLogge inner_text = utils.sanitize(inner_text) # If we don't have an inner text OR an attachment, we cancel because the hook can't send that - if not inner_text and not message.attachments: + # Strip so it counts a string of solely spaces as blank too + if not inner_text.strip() and not message.attachments: return False # So, we now have enough information to successfully proxy a message From f4512fc7e95de6da0eee705800a7ba124f72c2dc Mon Sep 17 00:00:00 2001 From: Ske Date: Thu, 8 Nov 2018 16:25:58 +0100 Subject: [PATCH 9/9] Various bug fixes --- .../bot/commands/message_commands.py | 5 +-- src/pluralkit/bot/commands/misc_commands.py | 2 +- src/pluralkit/bot/embeds.py | 42 ++++++++++--------- src/pluralkit/bot/proxy.py | 2 +- 4 files changed, 26 insertions(+), 25 deletions(-) diff --git a/src/pluralkit/bot/commands/message_commands.py b/src/pluralkit/bot/commands/message_commands.py index 7d474d8e..0fa1aafd 100644 --- a/src/pluralkit/bot/commands/message_commands.py +++ b/src/pluralkit/bot/commands/message_commands.py @@ -36,7 +36,7 @@ async def message_info(ctx: CommandContext): original_sender = None embed = discord.Embed() - embed.timestamp = discord.utils.snowflake_time(str(mid)) + embed.timestamp = discord.utils.snowflake_time(mid) embed.colour = discord.Colour.blue() if message.system_name: @@ -55,8 +55,7 @@ async def message_info(ctx: CommandContext): embed.add_field(name="Sent by", value=sender_name) message_content = await get_message_contents(ctx.client, message.channel, message.mid) - if message_content: - embed.description = message_content + embed.description = message_content or "(unknown, message deleted)" embed.set_author(name=message.name, icon_url=message.avatar_url or discord.Embed.Empty) diff --git a/src/pluralkit/bot/commands/misc_commands.py b/src/pluralkit/bot/commands/misc_commands.py index fde72a0a..803532c4 100644 --- a/src/pluralkit/bot/commands/misc_commands.py +++ b/src/pluralkit/bot/commands/misc_commands.py @@ -88,4 +88,4 @@ async def export(ctx: CommandContext): } f = io.BytesIO(json.dumps(data).encode("utf-8")) - await ctx.client.send_file(ctx.message.channel, f, filename="system.json") + await ctx.message.channel.send(content="Here you go!", file=discord.File(fp=f, filename="system.json")) diff --git a/src/pluralkit/bot/embeds.py b/src/pluralkit/bot/embeds.py index ec5bd038..4a5e1aa1 100644 --- a/src/pluralkit/bot/embeds.py +++ b/src/pluralkit/bot/embeds.py @@ -82,29 +82,31 @@ async def system_card(conn, client: discord.Client, system: System) -> discord.E value=system.description, inline=False) # Get names of all members - member_texts = [] - for member in await system.get_members(conn): - member_texts.append("{} (`{}`)".format(escape(member.name), member.hid)) + all_members = await system.get_members(conn) + if all_members: + member_texts = [] + for member in all_members: + member_texts.append("{} (`{}`)".format(escape(member.name), member.hid)) - # Interim solution for pagination of large systems - # Previously a lot of systems would hit the 1024 character limit and thus break the message - # This splits large system lists into multiple embed fields - # The 6000 character total limit will still apply here but this sort of pushes the problem until I find a better fix - pages = [""] - for member in member_texts: - last_page = pages[-1] - new_page = last_page + "\n" + member + # Interim solution for pagination of large systems + # Previously a lot of systems would hit the 1024 character limit and thus break the message + # This splits large system lists into multiple embed fields + # The 6000 character total limit will still apply here but this sort of pushes the problem until I find a better fix + pages = [""] + for member in member_texts: + last_page = pages[-1] + new_page = last_page + "\n" + member - if len(new_page) >= 1024: - pages.append(member) - else: - pages[-1] = new_page + if len(new_page) >= 1024: + pages.append(member) + else: + pages[-1] = new_page - for index, page in enumerate(pages): - field_name = "Members" - if index >= 1: - field_name = "Members (part {})".format(index + 1) - card.add_field(name=field_name, value=page, inline=False) + for index, page in enumerate(pages): + field_name = "Members" + if index >= 1: + field_name = "Members (part {})".format(index + 1) + card.add_field(name=field_name, value=page, inline=False) card.set_footer(text="System ID: {}".format(system.hid)) return card diff --git a/src/pluralkit/bot/proxy.py b/src/pluralkit/bot/proxy.py index 078f250f..ae4a7467 100644 --- a/src/pluralkit/bot/proxy.py +++ b/src/pluralkit/bot/proxy.py @@ -107,7 +107,7 @@ async def do_proxy_message(conn, original_message: discord.Message, proxy_member try: sent_message = await webhook.send( content=inner_text, - username=proxy_member.name, + username="{} {}".format(proxy_member.name, proxy_member.tag), avatar_url=proxy_member.avatar_url, file=await make_attachment_file(original_message), wait=True