Add support for image proxying

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

bot/pluralkit/proxy.py

@@ -1,4 +1,5 @@
 import os
+import json
 import time
 
 import aiohttp
@@ -19,6 +20,9 @@ async def log_message(original_message, hook_message, member, log_channel):
     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)
 
 async def get_webhook(conn, channel):
@@ -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