rc/main.py

#!/usr/bin/env python3
import pickle
import json
from typing import Callable, Awaitable
from dataclasses import dataclass, field
import discord
from discord.ext import commands
import traceback
from PIL import Image, ImageDraw
from io import BytesIO

bot = commands.Bot(command_prefix="\0", intents=discord.Intents.all())


@dataclass
class ReactionRole:
    emoji: str
    role_id: int
    description: str

    def __str__(self):
        return f"{self.emoji} - {self.description}"


@dataclass
class TrackedMessage:
    guild_id: int
    channel_id: int
    message_id: int
    header: str
    own_message: bool
    reactions: list[ReactionRole] = field(default_factory=list)

    def __str__(self):
        return f"message at https://discord.com/channels/{self.guild_id}/{self.channel_id}/{self.message_id} with {len(self.reactions)} reaction roles"

    def __repr__(self):
        return (self.header or " ") + '\n' + ("\n".join(
            str(rr) for rr in self.reactions) or f"-# waiting for reactions, message id is {self.message_id}")


class Database:
    def __init__(self, d: str):
        self.__tracked_messages: dict[int, TrackedMessage] = {}
        self.__file: str = d
        self.__load()

    def __load(self) -> None:
        try:
            self.__tracked_messages = pickle.load(open(self.__file, "rb"))
        except FileNotFoundError:
            self.__tracked_messages = {}

    def __save(self) -> None:
        pickle.dump(self.__tracked_messages, open(self.__file, "wb"))

    def track_message(self, guild_id: int, channel_id: int, message_id: int, header: str, own_message: bool) -> bool:
        if message_id in self.__tracked_messages:
            return True
        self.__tracked_messages[message_id] = TrackedMessage(guild_id, channel_id, message_id, header, own_message, [])
        self.__save()
        return False

    def edit_header(self, message_id: int, header: str) -> None:
        self.__tracked_messages[message_id].header = header

    def untrack_message(self, message_id: int):
        del self.__tracked_messages[message_id]
        self.__save()

    def add_reaction_role(self, message_id: int, emoji: str, role_id: int, description: str) -> bool:
        if message_id not in self.__tracked_messages:
            return True

        if any(rr.emoji == emoji or rr.role_id == role_id for rr in self.__tracked_messages[message_id].reactions):
            return True

        self.__tracked_messages[message_id].reactions.append(ReactionRole(emoji, role_id, description))
        self.__save()
        return False

    def edit_reaction_role_description(self, message_id: int, emoji: str, description: str) -> str | None:
        if message_id not in self.__tracked_messages:
            return None

        for rr in self.__tracked_messages[message_id].reactions:
            if rr.emoji == emoji:
                old_desc = rr.description
                rr.description = description
                self.__save()
                return old_desc

        return None

    def remove_reaction_role(self, message_id: int, emoji: str) -> None:
        self.__tracked_messages[message_id].reactions = [x for x in self.__tracked_messages[message_id].reactions if
                                                         x.emoji != emoji]
        self.__save()

    def embed_message(self, message_id: int) -> discord.Embed:
        m = self.__tracked_messages[message_id]
        return discord.Embed(
            title=f"message with {len(m.reactions)} reactions",
            description="\n".join(f"{x.emoji}" + (f" ({x.emoji.removesuffix('>').split(':')[2]})" if ':' in x.emoji else "") + f" - <@&{x.role_id}>" for x in m.reactions),
            color=discord.Color(0xa2d2df)
        )

    def format_message(self, message_id: int) -> str:
        return repr(self.__tracked_messages[message_id])

    def is_tracked_already(self, message_id: int):
        return message_id in self.__tracked_messages

    def is_my_message(self, message_id: int):
        return self.__tracked_messages[message_id].own_message

    async def get_message(self, message_id: int) -> discord.Message | None:
        if message_id not in self.__tracked_messages:
            return None
        tm = self.__tracked_messages[message_id]
        channel = bot.get_channel(tm.channel_id)
        if channel is None:
            return channel
        message = await channel.fetch_message(tm.message_id)
        return message

    def get_reaction_role(self, message_id: int, emoji: str) -> ReactionRole | None:
        if message_id not in self.__tracked_messages:
            return None
        for rr in self.__tracked_messages[message_id].reactions:
            if rr.emoji == emoji:
                return rr
        return None


class Config:
    def __init__(self, d: str):
        self.__file = d
        self.__load()

    def __load(self):
        try:
            c: dict[str, str | list[int]] = json.load(open(self.__file))
            self.__prefix: str = c["prefix"]
            self.__trusted: list[int] = c["trusted"]
            self.__token: str = c["token"]
        except FileNotFoundError:
            json.dumps({
                "prefix": "hey rc ",
                "trusted": [],
                "token": "https://youtu.be/dQw4w9WgXcQ"
            }, open(self.__file, "w"), indent=4)
            raise FileNotFoundError("Config not found")

    @property
    def prefix(self):
        return self.__prefix

    @property
    def trusted(self):
        return self.__trusted

    def run_bot(self):
        bot.run(self.__token)


class EmojiStorage:
    def __init__(self, d: str):
        self.__emoji_storage: dict[int, dict[str, int | list[str]]] = {}
        self.db_file: str = d
        self.__load()

    def __load(self) -> None:
        try:
            self.__emoji_storage = pickle.load(open(self.db_file, "rb"))
        except (FileNotFoundError, EOFError):
            self.__emoji_storage = {}

    def __save(self) -> None:
        pickle.dump(self.__emoji_storage, open(self.db_file, "wb"))

    def add_emoji_storage(self, server_id: int) -> bool:
        if server_id in self.__emoji_storage:
            return True

        self.__emoji_storage[server_id] = {"emoji_message_channel": 0, "emoji_message": 0, "emojis": []}
        self.__save()
        return False

    def set_emoji_message(self, server_id: int, channel_id: int, message_id: int):
        self.__emoji_storage[server_id]["emoji_message_channel"] = channel_id
        self.__emoji_storage[server_id]["emoji_message"] = message_id
        self.__save()

    def remove_emoji_storage(self, server_id: int) -> bool:
        if server_id not in self.__emoji_storage:
            return True

        del self.__emoji_storage[server_id]
        self.__save()
        return False

    async def __add_emoji(self, server_id: int, emoji: str):
        self.__emoji_storage[server_id]["emojis"].append(emoji)
        self.__save()
        channel = bot.get_channel(self.__emoji_storage[server_id]["emoji_message_channel"])
        message = await channel.fetch_message(self.__emoji_storage[server_id]["emoji_message"])
        await message.edit(content=''.join(self.__emoji_storage[server_id]["emojis"] or "-# dust"))

    async def try_add_emoji(self, name: str, image: bytes) -> discord.Emoji | None:
        for x in self.__emoji_storage:
            g = bot.get_guild(x)
            if len(g.emojis) == g.emoji_limit:
                continue
            emoji = await g.create_custom_emoji(name=name, image=image)
            await self.__add_emoji(x, str(emoji))
            return emoji
        return None


CONFIG: Config = Config("config.json")
DATABASE: Database = Database("db.dat")
EMOJI_STORAGE: EmojiStorage = EmojiStorage("emojis.dat")


rc_func = Callable[[list[str], discord.Message], Awaitable[tuple[str | discord.Embed, bool]]]
registry: list[tuple[str, rc_func, int]] = []


def register(name: str, args_amount: int):
    def decorator(func: rc_func):
        registry.append((name, func, args_amount))
        return func
    return decorator


@register("show", 1)
async def show(args: list[str], message: discord.Message) -> tuple[str | discord.Embed, bool]:
    if len(args) != 1:
        return "i dont understand your command", True

    if not args[0].isdigit():
        return "not a valid message id", True

    if not DATABASE.is_tracked_already(int(args[0])):
        return "not tracked", True

    return DATABASE.embed_message(int(args[0])), False


@register("send", 2)
async def send(args: list[str], message: discord.Message) -> tuple[str, bool]:
    if not message.author.guild_permissions.manage_guild:
        return "no permission", True

    if not 1 <= len(args) <= 2:
        return "i dont understand your command", True

    if not args[0].isdigit():
        return "not a valid channel id", True

    channel = bot.get_channel(int(args[0]))
    if channel is None:
        return "channel not found", True

    try:
        sent_message = await channel.send("hold on")
    except discord.Forbidden:
        return "no permission", True

    DATABASE.track_message(sent_message.guild.id, sent_message.channel.id, sent_message.id,
                           args[1] if len(args) == 2 else "", True)
    await sent_message.edit(content=DATABASE.format_message(sent_message.id))
    return "<:thumbsup:1486044381784834168>", True


@register("track", 2)
async def track(args: list[str], message: discord.Message) -> tuple[str, bool]:
    if not message.author.guild_permissions.manage_guild:
        return "no permission", True

    if len(args) != 2:
        return "i dont understand your command", True

    if not args[0].isdigit():
        return "not a valid channel id", True

    if not args[1].isdigit():
        return "not a valid message id", True

    if DATABASE.is_tracked_already(int(args[1])):
        return "already tracked", True

    channel = bot.get_channel(int(args[0]))
    if channel is None:
        return "channel not found", True

    try:
        tracked_msg = await channel.fetch_message(int(args[1]))
        DATABASE.track_message(tracked_msg.guild.id, tracked_msg.channel.id, tracked_msg.id, "", False)
        return f"successfully tracked message {args[1]}", True
    except discord.NotFound:
        return "message not found", True
    except discord.Forbidden:
        return "no permission to access that message", True


@register("add", 4)
async def add(args: list[str], message: discord.Message) -> tuple[str, bool]:
    if not message.author.guild_permissions.manage_guild:
        return "no permission", True

    if not 3 <= len(args) <= 4:
        return "i dont understand your command", True

    if not args[0].isdigit():
        return "not a valid message id", True

    if not DATABASE.is_tracked_already(int(args[0])):
        return "that message is not tracked. use `track` or `send` command first", True

    tracked_msg = await DATABASE.get_message(int(args[0]))
    if tracked_msg is None:
        return "could not fetch the message", True

    my = DATABASE.is_my_message(int(args[0]))

    if my and len(args) != 4:
        return "that's MY message and i want a description for it", True

    role = discord.utils.get(tracked_msg.guild.roles, name=args[2])
    if role is None:
        if args[2].isdigit():
            role = tracked_msg.guild.get_role(int(args[2]))
        else:
            return "not a role", True

    if args[1].isdigit():
        args[1] = bot.get_emoji(int(args[1]))
        if args[1] is None:
            return "emoji not found", True
        args[1] = str(args[1])

    if DATABASE.add_reaction_role(int(args[0]), args[1], role.id, args[3] if len(args) == 4 else ""):
        return "that emoji already has a reaction role on this message", True

    try:
        await tracked_msg.add_reaction(args[1])
    except discord.NotFound:
        return "invalid emoji", True
    except discord.Forbidden:
        return "missing permissions to add reactions", True
    except discord.HTTPException as e:
        return f"failed to add reaction: {e}", True

    if my:
        await tracked_msg.edit(content=DATABASE.format_message(int(args[0])))

    return f"successfully added reaction role: {args[1]} -> {role.name}" + (f" ({args[3]})" if len(args) == 4 else ""), True


@register("remove", 2)
async def remove(args: list[str], message: discord.Message) -> tuple[str, bool]:
    if not message.author.guild_permissions.manage_guild:
        return "no permission", True

    if len(args) != 2:
        return "i dont understand your command", True

    if not args[0].isdigit():
        return "not a valid message id", True

    if not DATABASE.is_tracked_already(int(args[0])):
        return "that message is not tracked", True

    if args[1].isdigit():
        args[1] = bot.get_emoji(int(args[1]))
        if args[1] is None:
            return "emoji not found", True
        args[1] = str(args[1])

    reaction_role = DATABASE.get_reaction_role(int(args[0]), args[1])
    if reaction_role is None:
        return "no reaction role found with that emoji on this message", True

    tracked_msg = await DATABASE.get_message(int(args[0]))
    if tracked_msg is None:
        return "could not fetch the message", True

    try:
        await tracked_msg.clear_reaction(args[1])
    except discord.Forbidden:
        return "missing permissions to remove reactions", True
    except discord.NotFound:
        pass
    except discord.HTTPException as e:
        return f"failed to remove reaction: {e}", True

    DATABASE.remove_reaction_role(int(args[0]), args[1])

    if DATABASE.is_my_message(int(args[0])):
        await tracked_msg.edit(content=DATABASE.format_message(int(args[0])))

    return f"successfully removed reaction role: {args[1]}", True


@register("delete", 1)
async def delete(args: list[str], message: discord.Message) -> tuple[str, bool]:
    if not message.author.guild_permissions.manage_guild:
        return "no permission", True

    if len(args) != 1:
        return "i dont understand your command", True

    if not args[0].isdigit():
        return "not a valid message id", True

    if not DATABASE.is_tracked_already(int(args[0])):
        return "that message is not tracked", True

    if not DATABASE.is_my_message(int(args[0])):
        return "i can only delete messages that i own (messages created with the `send` command)", True

    tracked_msg = await DATABASE.get_message(int(args[0]))

    if tracked_msg:
        try:
            await tracked_msg.delete()
        except discord.Forbidden:
            return "missing permissions to delete this message", True
        except discord.NotFound:
            pass
        except discord.HTTPException as e:
            return f"failed to delete message: {e}", True

    DATABASE.untrack_message(int(args[0]))

    return f"successfully deleted tracked message {int(args[0])}", True


@register("add this server to emoji storage", 0)
async def add_this_server_to_emoji_storage(args: list[str], message: discord.Message) -> tuple[str, bool]:
    if message.author.id not in CONFIG.trusted:
        return "no", True

    if EMOJI_STORAGE.add_emoji_storage(message.guild.id):
        return "already added", True

    sent_message = await message.channel.send("ok")
    EMOJI_STORAGE.set_emoji_message(sent_message.guild.id, sent_message.channel.id, sent_message.id)
    return "should've worked", True


@register("addemoji", 1)
async def addemoji(args: list[str], message: discord.Message) -> tuple[str, bool]:
    if message.guild.id != 1480697029780045875:
        return "this command cannot be ran in this server", True

    if not message.author.guild_permissions.manage_guild:
        return "no permission", True

    if len(args) != 1:
        return "i dont understand this command", True

    if len(message.attachments) != 1:
        return "nothing to add", True

    attachment = message.attachments[0]
    b: bytes = await attachment.read()
    emoji = await EMOJI_STORAGE.try_add_emoji(args[0], b)
    if emoji is None:
        return "failed to add emoji", True

    if args[0] == "arch":
        await message.channel.send("should've alpined")
    return f"added {emoji} with id {emoji.id}", False


@register("addcolor", 1)
async def addcolor(args: list[str], message: discord.Message) -> tuple[str, bool]:
    if message.guild.id != 1480697029780045875:
        return "this command cannot be ran in this server", True

    if not message.author.guild_permissions.manage_guild:
        return "no permission", True

    if len(args) != 1:
        return "i dont understand this command", True

    image = Image.new("RGBA", (128, 128), (0, 0, 0, 0))
    draw = ImageDraw.Draw(image)
    draw.rounded_rectangle([0, 0, 128, 128], radius=32, fill=args[0])
    buffer = BytesIO()
    image.save(buffer, format="PNG")
    image_bytes = buffer.getvalue()

    emoji = await EMOJI_STORAGE.try_add_emoji(args[0].removeprefix('#'), image_bytes)
    if emoji is None:
        return "failed to add emoji", True

    return f"added {emoji} with id {emoji.id}", False


@register("editdesc", 3)
async def editdesc(args: list[str], message: discord.Message) -> tuple[str, bool]:
    if not message.author.guild_permissions.manage_guild:
        return "no permission", True

    if len(args) != 3:
        return "i dont understand your command", True

    if not args[0].isdigit():
        return "not a valid message id", True

    if not DATABASE.is_tracked_already(int(args[0])):
        return "that message is not tracked. use `track` or `send` command first", True

    tracked_msg = await DATABASE.get_message(int(args[0]))
    if tracked_msg is None:
        return "could not fetch the message", True

    my = DATABASE.is_my_message(int(args[0]))
    if not my:
        return "can't help it it's not my message", True

    if args[1].isdigit():
        args[1] = bot.get_emoji(int(args[1]))
        if args[1] is None:
            return "emoji not found", True
        args[1] = str(args[1])

    result = DATABASE.edit_reaction_role_description(int(args[0]), args[1], args[2])
    if result is None:
        return "reaction role not found", True
    await tracked_msg.edit(content=DATABASE.format_message(int(args[0])))
    return f"changed reaction role description from {result} to {args[2]}", True


@register("edithdr", 2)
async def edithdr(args: list[str], message: discord.Message) -> tuple[str, bool]:
    if not message.author.guild_permissions.manage_guild:
        return "no permission", True

    if len(args) != 2:
        return "i dont understand your command", True

    if not args[0].isdigit():
        return "not a valid message id", True

    if not DATABASE.is_tracked_already(int(args[0])):
        return "that message is not tracked. use `track` or `send` command first", True

    tracked_msg = await DATABASE.get_message(int(args[0]))
    if tracked_msg is None:
        return "could not fetch the message", True

    my = DATABASE.is_my_message(int(args[0]))
    if not my:
        return "can't help it it's not my message", True

    DATABASE.edit_header(int(args[0]), args[1])
    await tracked_msg.edit(content=DATABASE.format_message(int(args[0])))
    return "i have done the thing", True


@bot.event
async def on_ready():
    print(f"{bot.user} is ready")
    await bot.change_presence(activity=discord.CustomActivity("gaming"))
    await bot.tree.sync()


async def handle_command(msg: str, message: discord.Message) -> tuple[str | discord.Embed, bool]:
    for reg in registry:
        if msg.lower().startswith(reg[0] + ' '):
            return await (reg[1])(
                msg.removeprefix(reg[0] + ' ').split(' ', max(0, reg[2] - 1)) if msg != reg[0] else [],
                message
            )

    return "", True


@bot.event
async def on_message(message: discord.Message):
    if message.content.lower() == "rc":
        await message.reply("reaction bot is here")
        return
        
    try:
        msg = message.content
        if not msg.lower().startswith(CONFIG.prefix):
            return

        result, d = await handle_command(msg.removeprefix(CONFIG.prefix), message)

        if not result:
            return

        if isinstance(result, str):
            if d:
                await message.reply(result, delete_after=60)
            else:
                await message.reply(result)
        elif isinstance(result, discord.Embed):
            await message.reply(embed=result)
    except Exception as e:
        await message.reply("```py\n" + '\n'.join(traceback.format_exception(e)) + "\n```")


@bot.event
async def on_raw_reaction_add(payload: discord.RawReactionActionEvent):
    if payload.user_id == bot.user.id:
        return

    reaction_role = DATABASE.get_reaction_role(payload.message_id, str(payload.emoji))
    if reaction_role is None:
        return

    guild = bot.get_guild(payload.guild_id)
    if guild is None:
        return

    member = guild.get_member(payload.user_id)
    if member is None:
        return

    role = guild.get_role(reaction_role.role_id)
    if role is None:
        return

    try:
        await member.add_roles(role, reason="reaction role")
    except discord.Forbidden:
        print(f"missing permissions to add role {role.name} to {member.name}")
    except discord.HTTPException as e:
        print(f"failed to add role: {e}")


@bot.event
async def on_raw_reaction_remove(payload: discord.RawReactionActionEvent):
    if payload.user_id == bot.user.id:
        return

    reaction_role = DATABASE.get_reaction_role(payload.message_id, str(payload.emoji))
    if reaction_role is None:
        return

    guild = bot.get_guild(payload.guild_id)
    if guild is None:
        return

    member = guild.get_member(payload.user_id)
    if member is None:
        return

    role = guild.get_role(reaction_role.role_id)
    if role is None:
        return

    try:
        await member.remove_roles(role, reason="reaction role")
    except discord.Forbidden:
        print(f"missing permissions to remove role {role.name} from {member.name}")
    except discord.HTTPException as e:
        print(f"failed to remove role: {e}")


@bot.tree.command(name="help", description="this description is here to help you use your help hope it helps")
async def help_slash_command(interaction: discord.Interaction):
    await interaction.response.send_message(embed=discord.Embed(
        title="reaction bot",
        description=rf"""made by tema5002
hosted by `[insert your name here at line {__import__('inspect').currentframe().f_lineno} at {__import__('os').path.abspath(__import__('os').path.abspath(__file__))}]`
https://gitea.codersquack.nl/tema5002/rc.git

"rc" stands for reaction clanker

Usage:
- {CONFIG.prefix}show message-id: Shows info about the message
- {CONFIG.prefix}send channel-id header (optional): Sends a message with `header` as it's top contents to be used for reaction roles
- {CONFIG.prefix}track channel-id message-id: Track an existing message for reaction roles
- {CONFIG.prefix}add message-id emoji/emoji-id role-id/role-name description: Adds a reaction role to message with `message-id` displayed as `emoji` `description` that gives you `role`
- {CONFIG.prefix}remove message-id emoji: Removes a reaction role from a tracked message
- {CONFIG.prefix}delete message-id: Deletes a tracked message that the bot owns
- {CONFIG.prefix}addemoji emoji-name + attachment: Adds emoji to the emoji storage. Only available for the "fuck alpinedevs,systemd,... server"
- {CONFIG.prefix}addcolor #RRGGBB: Same as addemoji but for color emojis
- {CONFIG.prefix}editdesc message-id emoji header: Edit description for a reaction roles emoji
- {CONFIG.prefix}edithdr message-id header: Edit header for a reaction roles message""",
        color=discord.Color(0xa2d2df)
    ))


CONFIG.run_bot()