muzovkant-cogdump/cogs/role_manager/role_manager.py

import discord
from discord.ext import commands
import logging
import asyncio
import json
import os

logger = logging.getLogger(__name__)

class RoleManager(commands.Cog):
    def __init__(self, bot):
        self.bot = bot
        self.config_path = os.path.join(os.path.dirname(__file__), 'config.json')
        self.config = self.load_config()

    def load_config(self):
        with open(self.config_path, 'r', encoding='utf-8') as f:
            return json.load(f)

    def save_config(self):
        with open(self.config_path, 'w', encoding='utf-8') as f:
            json.dump(self.config, f, indent=4, ensure_ascii=False)

    async def cog_load(self):
        asyncio.create_task(self._startup())

    async def _startup(self):
        await self.bot.wait_until_ready()
        if self.config.get("MESSAGE_ID"):
            await self.check_and_sync_roles()

    async def handle_reaction(self, payload, add_role=True):
        if payload.message_id != self.config.get("MESSAGE_ID"):
            return

        emoji = str(payload.emoji)
        role_id = self.config.get("REACTION_ROLES", {}).get(emoji)
        if not role_id:
            return

        guild = self.bot.get_guild(payload.guild_id)
        member = guild.get_member(payload.user_id)
        if not member or member.bot:
            return

        role = guild.get_role(role_id)
        if not role:
            return

        try:
            if add_role:
                await member.add_roles(role)
            else:
                await member.remove_roles(role)
        except discord.Forbidden:
            logger.error(f"Missing permissions for role {role.name}")

    @commands.Cog.listener()
    async def on_raw_reaction_add(self, payload):
        await self.handle_reaction(payload, add_role=True)

    @commands.Cog.listener()
    async def on_raw_reaction_remove(self, payload):
        await self.handle_reaction(payload, add_role=False)

    @commands.command(name="setup_role_menu")
    @commands.has_permissions(administrator=True)
    async def setup_role_menu(self, ctx):
        chan_id = self.config.get("CHANNEL_ID")
        channel = self.bot.get_channel(chan_id) or ctx.channel

        content = self.config.get("ROLE_MESSAGE_TEXT", "React for roles!")
        message = await channel.send(content)

        self.config["MESSAGE_ID"] = message.id
        self.config["CHANNEL_ID"] = channel.id
        self.save_config()

        for emoji in self.config.get("REACTION_ROLES", {}):
            await message.add_reaction(emoji)

    @commands.command(name="add_role")
    @commands.has_permissions(administrator=True)
    async def add_role(self, ctx, emoji: str, role: discord.Role):
        if "REACTION_ROLES" not in self.config:
            self.config["REACTION_ROLES"] = {}

        self.config["REACTION_ROLES"][emoji] = role.id
        self.save_config()
        if self.config.get("MESSAGE_ID"):
            try:
                chan = self.bot.get_channel(self.config["CHANNEL_ID"])
                msg = await chan.fetch_message(self.config["MESSAGE_ID"])
                await msg.add_reaction(emoji)
            except Exception as e:
                logger.error("failed to add reaction to menu: %s", e)

    @commands.command(name="remove_role")
    @commands.has_permissions(administrator=True)
    async def remove_role(self, ctx, emoji: str):
        if emoji in self.config.get("REACTION_ROLES", {}):
            del self.config["REACTION_ROLES"][emoji]
            self.save_config()

    @commands.command(name="edit_menu_text")
    @commands.has_permissions(administrator=True)
    async def edit_menu_text(self, ctx, *, new_text: str):
        self.config["ROLE_MESSAGE_TEXT"] = new_text
        self.save_config()
        if self.config.get("MESSAGE_ID"):
            try:
                chan = self.bot.get_channel(self.config["CHANNEL_ID"])
                msg = await chan.fetch_message(self.config["MESSAGE_ID"])
                await msg.edit(content=new_text)
                logger.info("menu text updated")
            except discord.NotFound:
                logger.error("could not find message when updating role message")
            except discord.Forbidden:
                logger.error("no permission err when updating role message")
            except discord.HTTPException as e:
                logger.error("http exception when updating role message: %s", e)
        else:
            logger.warning("text saved but no message id found, cannot update role menu text")

    async def check_and_sync_roles(self):
        try:
            chan = await self.bot.fetch_channel(self.config["CHANNEL_ID"])
            msg = await chan.fetch_message(self.config["MESSAGE_ID"])

            reacted_users = {}
            for reaction in msg.reactions:
                emoji = str(reaction.emoji)
                if emoji in self.config.get("REACTION_ROLES", {}):
                    reacted_users[emoji] = set()
                    async for user in reaction.users():
                        if not user.bot:
                            reacted_users[emoji].add(user.id)

            for emoji, role_id in self.config.get("REACTION_ROLES", {}).items():
                role = msg.guild.get_role(role_id)
                if not role:
                    continue

                reactors = reacted_users.get(emoji, set())

                for member in msg.guild.members:
                    if member.bot:
                        continue
                    has_role = role in member.roles
                    has_reacted = member.id in reactors

                    if has_reacted and not has_role:
                        await member.add_roles(role)
                    elif not has_reacted and has_role:
                        await member.remove_roles(role)

        except Exception as e:
            logger.error("check_and_sync_roles failed: %s", e)

async def setup(bot):
    await bot.add_cog(RoleManager(bot))