PluralKit/src/pluralkit/bot/proxy.py

from io import BytesIO

import discord
import logging
import re
from typing import List, Optional

from pluralkit import db
from pluralkit.bot import utils, channel_logger
from pluralkit.bot.channel_logger import ChannelLogger

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)
    if not match:
        return message_text, ""

    # Return the text after the mentions, and the mentions themselves
    return message_text[match.span(0)[1]:].strip(), match.group(0)


def match_member_proxy_tags(member: db.ProxyMember, message_text: str):
    # Skip members with no defined proxy tags
    if not member.prefix and not member.suffix:
        return None

    # DB defines empty prefix/suffixes as None, replace with empty strings to prevent errors
    prefix = member.prefix or ""
    suffix = member.suffix or ""

    # 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))

    if message_text.startswith(member.prefix or "") and message_text.endswith(member.suffix or ""):
        prefix_length = len(prefix)
        suffix_length = len(suffix)

        # If suffix_length is 0, the last bit of the slice will be "-0", and the slice will fail
        if suffix_length > 0:
            inner_string = message_text[prefix_length:-suffix_length]
        else:
            inner_string = message_text[prefix_length:]

        # Add the mentions we stripped back
        inner_string = leading_mentions + inner_string
        return inner_string


def match_proxy_tags(members: List[db.ProxyMember], message_text: str):
    # Sort by specificity (members with both prefix and suffix go higher)
    # This will make sure more "precise" proxy tags get tried first
    members: List[db.ProxyMember] = sorted(members, key=lambda x: int(
        bool(x.prefix)) + int(bool(x.suffix)), reverse=True)

    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
            logger.debug("Matched member {} with inner text '{}'".format(member.hid, match))
            return member, match


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")
    created_webhook._adapter.store_user = created_webhook._adapter._store_user

    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

    first_attachment = message.attachments[0]

    # 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)


async def do_proxy_message(conn, original_message: discord.Message, proxy_member: db.ProxyMember,
                           inner_text: str, logger: ChannelLogger):
    # Send the message through the webhook
    webhook = await get_or_create_webhook_for_channel(conn, original_message.channel)

    try:
        sent_message = await webhook.send(
            content=inner_text,
            username="{} {}".format(proxy_member.name, proxy_member.tag),
            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,
                         original_message.author.id)

    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(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

    # 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

    member, inner_text = proxy_match

    # Sanitize inner text for @everyones and such
    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
    # 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
    async with conn.transaction():
        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)