PluralKit/bot/pluralkit/bot.py

from datetime import datetime
import logging
import json
import os
import time

import discord

logging.basicConfig(level=logging.INFO, format="[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s")
logging.getLogger("discord").setLevel(logging.INFO)
logging.getLogger("websockets").setLevel(logging.INFO)

logger = logging.getLogger("pluralkit.bot")
logger.setLevel(logging.INFO)

client = discord.Client()


@client.event
async def on_error(evt, *args, **kwargs):
    logger.exception(
        "Error while handling event {} with arguments {}:".format(evt, args))


@client.event
async def on_ready():
    # Print status info
    logger.info("Connected to Discord.")
    logger.info("Account: {}#{}".format(
        client.user.name, client.user.discriminator))
    logger.info("User ID: {}".format(client.user.id))


@client.event
async def on_message(message):
    # Ignore bot messages
    if message.author.bot:
        return

    # Split into args. shlex sucks so we don't bother with quotes
    args = message.content.split(" ")

    from pluralkit import proxy, utils, stats
    
    command_items = utils.command_map.items()
    command_items = sorted(command_items, key=lambda x: len(x[0]), reverse=True)
    
    prefix = "pk;"
    for command, (func, _, _, _) in command_items:
        if message.content.lower().startswith(prefix + command):
            args_str = message.content[len(prefix + command):].strip()
            args = args_str.split(" ")

            # Splitting on empty string yields one-element array, remove that
            if len(args) == 1 and not args[0]:
                args = []

            async with client.pool.acquire() as conn:
                time_before = time.perf_counter()
                await func(conn, message, args)
                time_after = time.perf_counter()

                # Report command time stats
                execution_time = time_after - time_before
                response_time = (datetime.now() - message.timestamp).total_seconds()
                await stats.report_command(command, execution_time, response_time)
                return

    # Try doing proxy parsing
    async with client.pool.acquire() as conn:
        await proxy.handle_proxying(conn, message)

@client.event
async def on_socket_raw_receive(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 client.pool.acquire() as conn:
                        from pluralkit import proxy
                        await proxy.handle_reaction(conn, user_id, message_id, emoji)
        except ValueError:
            pass

async def run():
    from pluralkit import db, stats
    try:
        logger.info("Connecting to database...")
        pool = await db.connect()

        logger.info("Attempting to create tables...")
        async with pool.acquire() as conn:
            await db.create_tables(conn)

        logger.info("Connecting to InfluxDB...")
        await stats.connect()

        client.pool = pool
        logger.info("Connecting to Discord...")
        await client.start(os.environ["TOKEN"])
    finally:
        logger.info("Logging out from Discord...")
        await client.logout()