rollbot4/drivers/discord_driver.py

from __future__ import annotations

import io
import asyncio
import logging.config
import os

import toml
import discord

import rollbot
from commands import config

logging.config.fileConfig("logging.conf", disable_existing_loggers=False)

with open(os.environ.get("SECRET_FILE", "secrets.toml"), "r") as sfile:
    secrets = toml.load(sfile)
database_file = os.environ.get("DATABASE_FILE", secrets["database_file"])


class DiscordBot(rollbot.Rollbot[discord.Message]):
    def __init__(self, client):
        super().__init__(config.extend(rollbot.CommandConfiguration(bangs=("!",))), database_file)
        self.discord_client = client
        client.event(self.on_message)

    def read_config(self, key: str) -> Any:
        cfg = secrets
        for part in key.split("."):
            cfg = cfg.get(part, None)
            if cfg is None:
                return None
        return cfg

    async def parse(self, msg: discord.Message) -> rollbot.Message:
        return rollbot.Message(
            origin_id="DISCORD",
            channel_id=str(msg.channel.id),
            sender_id=str(msg.author.id),
            timestamp=msg.created_at,
            origin_admin="RollbotAdmin" in [r.name for r in msg.author.roles],
            channel_admin=False, # TODO - implement this if discord allows it
            sender_name=msg.author.name,
            text=msg.content,
            # TODO might be nice to only read attachments once we're sure we need to
            attachments=[rollbot.Attachment(
                name=att.filename,
                body=await att.read(),
            ) for att in msg.attachments],
        )

    async def respond(self, response: rollbot.Response):
        if response.origin_id != "DISCORD":
            self.context.logger.error(f"Unable to respond to {response.origin_id}")
            return
        
        channel = await self.discord_client.fetch_channel(response.channel_id)

        content=response.text or ""
        # TODO this implementation is wrong - attaches files instead of embedding an image
        # I *think* the right call is to send multiple messages since there's only one embed each
        # but I want to do a little more research before implementing that
        files = []
        if response.attachments is not None:
            for att in response.attachments:
                if att.name == "image":
                    if isinstance(att.body, bytes):
                        files.append(discord.File(io.BytesIO(att.body)))
                    else:
                        content += "\n" + att.body

        await channel.send(content=content, files=files)


if __name__ == "__main__":
    loop = asyncio.get_event_loop()
    client = discord.Client(loop=loop)
    bot = DiscordBot(client)
    try:
        loop.run_until_complete(bot.on_startup())
        loop.run_until_complete(client.start(secrets["discord"]["token"]))
    except KeyboardInterrupt:
        loop.run_until_complete(client.close())
    finally:
        loop.run_until_complete(bot.on_shutdown())