rollbot4/drivers/discord_driver.py

from __future__ import annotations

import io
import asyncio
import logging.config
import os
from typing import Any

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"])
bangs = tuple(t for t in secrets["discord"].get("bangs", "!/"))


class DiscordBot(rollbot.Rollbot[discord.Message]):
    def __init__(self, client):
        super().__init__(config.extend(rollbot.CommandConfiguration(bangs=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:
        # TODO might be nice to only read attachments lazily

        attachments = [rollbot.Attachment(
            name=att.filename,
            body=await att.read(),
        ) for att in msg.attachments]

        if msg.reference is not None:
            channel = await self.discord_client.fetch_channel(msg.channel.id)
            attachments.append(rollbot.Attachment(
                name="reply",
                body=await channel.fetch_message(msg.reference.resolved.id),
            ))

        return rollbot.Message(
            origin_id="DISCORD",
            channel_id=str(msg.channel.id),
            sender_id=str(msg.author.id),
            message_id=str(msg.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,
            attachments=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)

        args = {}
        args["content"] = response.text or ""
        attachments = []
        files = []
        if response.attachments is not None:
            for att in response.attachments:
                if att.name == "image":
                    if isinstance(att.body, bytes):
                        embed = discord.Embed(description="Embedded Image")
                        file = discord.File(io.BytesIO(att.body), filename="image.png")
                        embed.set_image(url="attachment://image.png")
                        args["embed"] = embed
                        # TODO might eventually be nice to figure out a way of doing multiple embeds
                        files.append(file)
                    else:
                        args["content"] += "\n" + att.body
                elif att.name == "reply":
                    if att.body is None or not isinstance(att.body, str):
                        raise ValueError("Invalid reply body type, must be message ID")
                    args["reference"] = await channel.fetch_message(int(att.body))
                elif isinstance(att.body, discord.Attachment):
                    attachments.append(att.body)
        if len(attachments) > 0:
            args["attachments"] = attachments
        if len(files) > 0:
            args["files"] = files

        await channel.send(**args)


if __name__ == "__main__":
    loop = asyncio.get_event_loop()
    intents = discord.Intents.default()
    intents.message_content = True
    client = discord.Client(intents=intents, 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())