From fc93ebccbbb6b93164c4f6324ac4e77e31046ca0 Mon Sep 17 00:00:00 2001 From: Ske Date: Fri, 13 Jul 2018 21:11:31 +0200 Subject: [PATCH] Add support for image proxying --- bot/pluralkit/db.py | 2 +- bot/pluralkit/proxy.py | 54 +++++++++++++++++++++++++++++------------- 2 files changed, 39 insertions(+), 17 deletions(-) diff --git a/bot/pluralkit/db.py b/bot/pluralkit/db.py index 2f48d95e..2e60f34b 100644 --- a/bot/pluralkit/db.py +++ b/bot/pluralkit/db.py @@ -156,7 +156,7 @@ async def add_message(conn, message_id: str, channel_id: str, member_id: int, se @db_wrap async def get_members_by_account(conn, account_id: str): # Returns a "chimera" object - return await conn.fetch("select members.id, members.hid, members.prefix, members.suffix, members.name, members.avatar_url, systems.tag, systems.name as system_name, systems.hid as system_hid from systems, members, accounts where accounts.uid = $1 and systems.id = accounts.system and members.system = systems.id", int(account_id)) + return await conn.fetch("select members.id, members.hid, members.prefix, members.suffix, members.color, members.name, members.avatar_url, systems.tag, systems.name as system_name, systems.hid as system_hid from systems, members, accounts where accounts.uid = $1 and systems.id = accounts.system and members.system = systems.id", int(account_id)) @db_wrap diff --git a/bot/pluralkit/proxy.py b/bot/pluralkit/proxy.py index 4670dd70..68179d50 100644 --- a/bot/pluralkit/proxy.py +++ b/bot/pluralkit/proxy.py @@ -1,4 +1,5 @@ import os +import json import time import aiohttp @@ -18,6 +19,9 @@ async def log_message(original_message, hook_message, member, log_channel): embed.timestamp = hook_message.timestamp embed.set_author(name=author_name, icon_url=member["avatar_url"] or discord.Embed.Empty) embed.set_footer(text="System ID: {} | Member ID: {} | Sender: {}#{} | Message ID: {}".format(member["system_hid"], member["hid"], original_message.author.name, original_message.author.discriminator, hook_message.id)) + + if len(hook_message.attachments) > 0: + embed.set_image(url=hook_message.attachments[0]["url"]) await client.send_message(log_channel, embed=embed) @@ -43,20 +47,34 @@ async def get_webhook(conn, channel): return hook_row["webhook"], hook_row["token"] -async def send_hook_message(member, text, hook_id, hook_token): +async def send_hook_message(member, hook_id, hook_token, text=None, image_url=None): async with aiohttp.ClientSession() as session: - # Set up parameters - req_data = { - "username": "{} {}".format(member["name"], member["tag"] or "").strip(), - "avatar_url": member["avatar_url"], - "content": text + # Set up headers + req_headers = { + "Authorization": "Bot {}".format(os.environ["TOKEN"]) } - req_headers = {"Authorization": "Bot {}".format(os.environ["TOKEN"])} - # Send request - async with session.post("https://discordapp.com/api/v6/webhooks/{}/{}?wait=true".format(hook_id, hook_token), json=req_data, headers=req_headers) as resp: + # Set up parameters + # Use FormData because the API doesn't like JSON requests with file data + fd = aiohttp.FormData() + fd.add_field("username", "{} {}".format(member["name"], member["tag"] or "").strip()) + fd.add_field("avatar_url", member["avatar_url"]) + + if text: + fd.add_field("content", text) + + if image_url: + # Fetch the image URL and proxy it directly into the file data (async streaming!) + image_resp = await session.get(image_url) + fd.add_field("file", image_resp.data, content_type=image_resp.content_type, filename=image_resp.url.name) + + # Send the actual webhook request, and wait for a response + async with session.post("https://discordapp.com/api/v6/webhooks/{}/{}?wait=true".format(hook_id, hook_token), + data=fd, + headers=req_headers) as resp: if resp.status == 200: resp_data = await resp.json() + # Make a fake message object for passing on - this is slightly broken but works for most things return discord.Message(reactions=[], **resp_data) else: # Fake a Discord exception, also because #yolo @@ -64,16 +82,17 @@ async def send_hook_message(member, text, hook_id, hook_token): async def proxy_message(conn, member, trigger_message, inner): - logger.debug("Proxying message '{}' for member {}".format( - inner, member["hid"])) - # Delete the original message - await client.delete_message(trigger_message) + logger.debug("Proxying message '{}' for member {}".format(inner, member["hid"])) # Get the webhook details hook_id, hook_token = await get_webhook(conn, trigger_message.channel) - # And send the message - hook_message = await send_hook_message(member, inner, hook_id, hook_token) + # Get attachment image URL if present (only works for one...) + image_urls = [a["url"] for a in trigger_message.attachments if "url" in a] + image_url = image_urls[0] if len(image_urls) > 0 else None + + # Send the hook message + hook_message = await send_hook_message(member, hook_id, hook_token, text=inner, image_url=image_url) # Insert new message details into the DB await db.add_message(conn, message_id=hook_message.id, channel_id=trigger_message.channel.id, member_id=member["id"], sender_id=trigger_message.author.id) @@ -83,9 +102,12 @@ async def proxy_message(conn, member, trigger_message, inner): if server_info and server_info["log_channel"]: channel = trigger_message.server.get_channel(str(server_info["log_channel"])) if channel: - # Log the message + # Log the message to the log channel if present await log_message(trigger_message, hook_message, member, channel) + # Delete the original message + await client.delete_message(trigger_message) + async def handle_proxying(conn, message): # Big fat query to find every member associated with this account