import re from discord import Embed, Member, Guild from discord.errors import Forbidden from discord.ext import commands from discord.ext.commands import BadArgument import db from administrator.check import is_enabled from administrator.logger import logger from administrator.utils import event_is_enabled extension_name = "invite" logger = logger.getChild(extension_name) role_mention_re = re.compile(r"<@&[0-9]+>") channel_mention_re = re.compile(r"<#[0-9]+>") class Invite(commands.Cog): def __init__(self, bot: commands.Bot): self.bot = bot self.invites = {} self.bot.loop.create_task(self.update_invites()) def description(self): return "Get role from a special invite link" @commands.group("invite", pass_context=True) @is_enabled() @commands.guild_only() @commands.has_guild_permissions(administrator=True) async def invite(self, ctx: commands.Context): if ctx.invoked_subcommand is None: await ctx.invoke(self.invite_help) @invite.group("help", pass_context=True) async def invite_help(self, ctx: commands.Context): embed = Embed(title="Invite help") embed.add_field(name="invite create <#channel> <@role>", value="Create a invite link to a role", inline=False) embed.add_field(name="invite delete ", value="Remove a invite", inline=False) await ctx.send(embed=embed) @invite.group("create", pass_context=True) async def invite_add(self, ctx: commands.Context, channel: str, role: str): if not channel_mention_re.fullmatch(channel) or len(ctx.message.channel_mentions) != 1 or\ not role_mention_re.fullmatch(role) or len(ctx.message.role_mentions) != 1: raise BadArgument() inv = await ctx.message.channel_mentions[0].create_invite() s = db.Session() s.add(db.InviteRole(ctx.guild.id, inv.code, ctx.message.role_mentions[0].id)) s.commit() s.close() await ctx.send(f"Invite created: `{inv.url}`") @invite.group("delete", pass_context=True) async def invite_delete(self, ctx: commands.Context, code: str): inv = next(filter(lambda i: i.code == code, await ctx.guild.invites()), None) if not inv: raise BadArgument() s = db.Session() invite_role = s.query(db.InviteRole).get({"guild_id": ctx.guild.id, "invite_code": code}) if not invite_role: s.close() raise BadArgument() s.delete(invite_role) s.commit() s.close() await inv.delete() await ctx.message.add_reaction("\U0001f44d") async def update_invites(self): for g in self.bot.guilds: self.invites[g.id] = await g.invites() @commands.Cog.listener() async def on_ready(self): await self.update_invites() @commands.Cog.listener() async def on_member_join(self, member: Member): if not event_is_enabled(self.qualified_name, member.guild.id): return user_invites = await member.guild.invites() for i in self.invites[member.guild.id]: for ui in user_invites: if i.code == ui.code and i.uses < ui.uses: s = db.Session() invite_role = s.query(db.InviteRole).get({"guild_id": member.guild.id, "invite_code": i.code}) s.close() if invite_role: try: await member.add_roles(member.guild.get_role(invite_role.role_id)) except Forbidden: pass self.invites[member.guild.id] = user_invites @commands.Cog.listener() async def on_invite_create(self, invite): if not event_is_enabled(self.qualified_name, invite.guild.id): return self.invites[invite.guild.id] = await invite.guild.invites() @commands.Cog.listener() async def on_invite_delete(self, invite): s = db.Session() if not event_is_enabled(self.qualified_name, invite.guild.id, s): return invite_role = s.query(db.InviteRole).get({"guild_id": invite.guild.id, "invite_code": invite.code}) if invite_role: s.delete(invite_role) s.commit() s.close() self.invites[invite.guild.id] = await invite.guild.invites() @commands.Cog.listener() async def on_guild_join(self, guild: Guild): self.invites[guild.id] = await guild.invites() @commands.Cog.listener() async def on_guild_remove(self, guild: Guild): s = db.Session() for g in s.query(db.InviteRole).filter(db.InviteRole.guild_id == guild.id).all(): s.delete(g) s.commit() s.close() del self.invites[guild.id] def setup(bot): logger.info(f"Loading...") try: bot.add_cog(Invite(bot)) except Exception as e: logger.error(f"Error loading: {e}") else: logger.info(f"Load successful") def teardown(bot): logger.info(f"Unloading...") try: bot.remove_cog("Invite") except Exception as e: logger.error(f"Error unloading: {e}") else: logger.info(f"Unload successful")