Current use of shoving data in custom_id

stripped_down.py

#   Copyright 2023-present Michael Hall
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

from __future__ import annotations

import logging
import os
import re
from typing import Self

import base2048
import discord
import msgpack
from cryptography.exceptions import InvalidTag
from cryptography.hazmat.primitives.ciphers.aead import AESSIV

log = logging.getLogger(__name__)

def get_aessiv_key():  # real code uses systemd-creds here
    with open("aessiv.key", mode="rb") as f:
        return f.read()


def get_discord_token():  # real code uses systemd-creds here
    return os.getenv("ROLEBOT_DISCORD_TOKEN")


class RoleBot(discord.AutoShardedClient):

    def __init__(self, aessiv: AESSIV, *args, **kwargs) -> None:
        self.aessiv = aessiv
        self.interaction_regex = re.compile(r"^rr\d{2}:(.*)$", flags=re.DOTALL)
        super().__init__(*args, **kwargs)

    async def on_interaction(self: Self, interaction: discord.Interaction[Self]):
        if interaction.type is discord.InteractionType.component:
            if m := self.interaction_regex.match(interaction.data["custom_id"]):
                await self.handle_rules_for_interaction(interaction, m.group(1))

    async def handle_rules_for_interaction(self: Self, interaction: discord.Interaction[Self], encoded_rules: str):
        await interaction.response.defer(ephemeral=True)

        try:
             by = base2048.decode(encoded_rules)
        except Exception:
            log.exception("Could not decode expected valid custom id using base2048: %s", encoded_rules)
            return

        try:
            dec = self.aessiv.decrypt(by, [interaction.guild_id.to_bytes(64)])
        except InvalidTag:
            log.exception("Got a custom id that couldn't decrypt with aessiv key", exc_info=True)
            return
        
        try:
            rules_object: dict = msgpack.loads(dec, strict_map_key=False, use_list=False)
        except Exception:
            log.exception("Could not unpack after decryption passed: %s", dec)
            return

        member: discord.Member = interaction.user

        toggle = set(rules_object.get("^", ()))
        add = set(rules_object.get("|", ()))
        remove = set(rules_object.get("-", ()))

        role_ids = {r.id for r in member.roles}

        role_ids ^= toggle
        role_ids |= add
        role_ids -= remove

        await member.edit(roles=[discord.Object(rid) for rid in role_ids], reason="Used a role button.")

    async def create_role_menu(
        self: Self,
        channel: discord.TextChannel,
        content: str | None,
        embed: discord.Embed | None,
        label_action_pairs: list[tuple[str, dict]],
    ):
        
        if len(label_action_pairs) > 25:
            raise ValueError("Cannot provide that many buttons in a view")

        if not (content or embed):
            raise ValueError("Must provide an embed or message content.")
    
        guild_id_b = channel.guild.id.to_bytes(64)

        view = discord.ui.View(timeout=None)
        view.stop()

        for idx, (label, action_dict) in enumerate(label_action_pairs):
            data = msgpack.dumps(action_dict)
            enc = self.aessiv.encrypt(data, [guild_id_b])
            to_disc = base2048.encode(enc)
            custom_id = f"rr{idx:02}:{to_disc}"
            button = discord.ui.Button(label=label, custom_id=custom_id)
            view.add_item(button)

        await channel.send(content=content, embed=embed, view=view)


# I've left out how users interact with the bot to create a role menu
# for now because the current method sucks and needs re-doing.
# Current method only binds 1 action (Add/remove/toggle) to 1 role per button,
# while the structure here allows more.