from redbot.core.utils.chat_formatting import * from redbot.core import Config, checks, commands, bank from typing import Literal import discord import asyncio class PoorError(commands.CheckFailure): pass # 0 -> member, 1 -> guild NEW_RECEIPT_MESSAGE = "Hello {0.mention}! This is will be your receipt for commands you pay for in {1.name}.\nThe message below will auto update as you pay for commands. I will pin the message for you so you can refer back to this to track your spending." # 0 -> recepit number, 1 -> command name, 2->command cost, 3 -> currency name RECEIPT_MESSAGE = "{0}. {1}: {2} {3}\n" MAX_MSG_LEN = 2000 class CostManager(commands.Cog): """ Allows customizing costs for any commands loaded into read, and also set users who are exempt from costs on a per member or per role level. Hierarchy: user > role > guild_role """ def __init__(self, bot): self.bot = bot self.config = Config.get_conf(self, identifier=13291493293, force_registration=True) # commands: { # command_name: { # cost: int # user_ids: {user_id:cost} # role_ids: {role_id:cost} # cost can be 0 to make free # }, # ... # } default_guild = {"FREE_ROLES": [], "COMMANDS": {}} self.config.register_guild(**default_guild) self.config.register_member(receipt=0) self.bot.before_invoke(self.cost_checker) def cog_unload(self): self.bot.remove_before_invoke_hook(self.cost_checker) # permission hook checker for cost of command async def cost_checker(self, ctx): cost = await self.get_cost(ctx) if cost == 0: return None try: await bank.withdraw_credits(ctx.author, cost) await self.update_receipt(ctx, cost) except ValueError: raise PoorError(f"member: {ctx.author.id}, guild: {ctx.guild.id}") async def get_cost(self, ctx, member=None, command=None): """ Get cost of a command, respecting hierarchy """ if isinstance(ctx.channel, discord.DMChannel): return 0 guild = ctx.guild if not member: member = ctx.author member_roles = {r.id for r in member.roles if r.name != "@everyone"} if not command: command = ctx.command.name guild_data = await self.config.guild(guild).all() command_data = guild_data["COMMANDS"] # check if settings for command if command not in command_data.keys(): return 0 command_data = command_data[command] charged_roles = set(command_data.get("role_ids", {}).keys()) found_roles = charged_roles & member_roles cost = 0 # check user cost if str(member.id) in command_data.get("user_ids", {}).keys(): cost = command_data["user_ids"][str(member.id)] # check role cost, choose lowest cost if mutliple roles found. elif found_roles: cost = min([command_data["role_ids"][r] for r in found_roles]) else: # get normal cost for command and check guild free roles cost = command_data["cost"] found_roles = set(guild_data["FREE_ROLES"]) & member_roles if found_roles: cost = 0 return cost async def update_receipt(self, ctx, cost): """ Update user's receipt, or make a new one if its new receipt or old one is somehow missing. """ guild = ctx.guild member = ctx.author command = ctx.command.name receipt = await self.config.member(member).receipt() channel = member.dm_channel if not channel: await member.create_dm() channel = member.dm_channel msg_receipt = None if receipt > 0: try: msg_receipt = await channel.fetch_message(receipt) except (Forbidden, HTTPException): return except NotFound: pass if not msg_receipt: try: await channel.send(NEW_RECEIPT_MESSAGE.format(member, guild)) except: if receipt != -1: # send mention if first time getting this message await ctx.send( f"Hey {member.mention}, please turn on DMs from server members in your settings so I can send you receipts for purchase of commands.", allowed_mentions=discord.AllowedMentions.all(), ) await self.config.member(member).receipt.set(-1) return currency_name = await bank.get_currency_name(guild) msg_receipt = await channel.send(RECEIPT_MESSAGE.format(1, command, cost, currency_name)) await msg_receipt.pin() await self.config.member(member).receipt.set(msg_receipt.id) return # msg already found, so update it currency_name = await bank.get_currency_name(guild) content = msg_receipt.content.split("\n") # TODO: this is really messy but ill fix it... last_num = int(content[-1].split(".")[0]) content.append(RECEIPT_MESSAGE.format(last_num + 1, command, cost, currency_name)) while len("\n".join(content)) > MAX_MSG_LEN: del content[0] content = "\n".join(content) await msg_receipt.edit(content=content) async def clean_data(self, guild): async with self.config.guild(guild).FREE_ROLES() as free: roles = [guild.get_role(id) for id in free] free = [r.id for r in roles if r is not None] async with self.config.guild(guild).COMMANDS() as c: for command_name in list(c.keys()): data = c[command_name] for user_id in data.get("user_ids", {}).keys(): if not guild.get_member(int(user_id)): del c[command_name]["user_ids"][user_id] for role_id in data.get("role_ids", {}).keys(): if not guild.get_role(int(role_id)): del c[command_name]["role_ids"][role_id] # check validity of arguments def arg_check(self, cost: int, command: str): if not self.bot.get_command(command) or cost < 0: return False return True # format list of items @staticmethod def format_list(*items, join="and", delim=", "): if len(items) > 1: return (" %s " % join).join((delim.join(items[:-1]), items[-1])) elif items: return items[0] else: return "" @commands.group(name="costset", invoke_without_command=True) @commands.guild_only() @checks.admin() async def costset(self, ctx, cost: int, *, command_name: str = None): """ Sets and manage cost/bypasses of commands. """ if ctx.invoked_subcommand: return if not self.arg_check(cost, command_name): await ctx.send("Invalid cost or command!") return async with self.config.guild(ctx.guild).COMMANDS() as c: if not command_name in c.keys(): c[command_name] = {} c[command_name]["cost"] = cost await ctx.tick() @costset.command(name="role") async def cost_set_role(self, ctx, command_name: str, cost: int, *, role: discord.Role): """ Set cost of command for specific role. Set to 0 to make command free for role. """ if not self.arg_check(cost, command_name): await ctx.send("Invalid cost or command!") return async with self.config.guild(ctx.guild).COMMANDS() as c: if not command_name in c.keys(): c[command_name] = {} if not "role_ids" in c[command_name].keys(): c[command_name]["role_ids"] = {} c[command_name]["role_ids"][str(role.id)] = cost await ctx.tick() @costset.command(name="user") async def cost_set_user(self, ctx, command_name: str, cost: int, *, member: discord.Member): """ Set cost of command for specific user. Set to 0 to make command free for user. """ if not self.arg_check(cost, command_name): await ctx.send("Invalid cost or command!") return async with self.config.guild(ctx.guild).COMMANDS() as c: if not command_name in c.keys(): c[command_name] = {} if not "user_ids" in c[command_name].keys(): c[command_name]["user_ids"] = {} c[command_name]["user_ids"][str(member.id)] = cost await ctx.tick() @costset.command(name="clear") async def cost_set_clear(self, ctx): """ Clear missing roles/members in cost config. """ await self.clean_data(ctx.guild) await ctx.tick() @costset.command(name="owner-clear") @checks.is_owner() async def cost_set_owner_clear(self, ctx): """ Clear missing roles/members in every cost config. """ for guild in self.bot.guilds: await self.clean_data(guild) await ctx.tick() @costset.command(name="free-roles") async def cost_set_free_roles(self, ctx, *, role_list: str = None): """ Set roles who can use all commands for free in server. **Note**: This only applies to commands whose cost was set with this cog. Role list should be a list of one or more **role names or ids** seperated by commas. Roles in role list will be removed if already in the free role list, or added if they are not. Role names are case sensitive! Don't pass a role list to see the current roles """ if not role_list: curr = await self.config.guild(ctx.guild).FREE_ROLES() if not curr: await ctx.send("No roles defined.") else: curr = [ctx.guild.get_role(role_id) for role_id in curr] not_found = len([r for r in curr if r is None]) curr = [r.name for r in curr if curr is not None] if not_found: await ctx.send( f"{not_found} roles weren't found, please run {ctx.prefix}costset clear to remove these roles.\nFree Roles: {self.format_list(*curr)}" ) else: await ctx.send(f"Free Roles: {self.format_list(*curr)}") return role_list = role_list.strip().split(",") role_list = [r.strip() for r in role_list] not_found = set() found = set() added = set() removed = set() for role_name in role_list: role = discord.utils.find(lambda r: r.name == role_name, ctx.guild.roles) # if couldnt find by role name, try to find by role id if role is None: role = discord.utils.find(lambda r: r.id == role_name, ctx.guild.roles) if role is None: not_found.add(role_name) continue found.add(role) if not_found: await ctx.send( warning("These roles weren't found, please try again: {}".format(self.format_list(*not_found))) ) return async with self.config.guild(ctx.guild).FREE_ROLES() as free: for role in found: if role.id in free: free.remove(role.id) removed.add(role.name) else: free.append(role.id) added.add(role.name) msg = "" if added: msg += "Added: {}\n".format(self.format_list(*added)) if removed: msg += "Removed: {}".format(self.format_list(*removed)) await ctx.send(msg) @costset.command(name="list") async def cost_set_list(self, ctx): """ List current cost settings for the guild """ guild = ctx.guild guild_data = await self.config.guild(guild).all() free_roles = [guild.get_role(r) for r in guild_data["FREE_ROLES"]] free_roles = self.format_list(*[r.name for r in free_roles if r is not None]) commands = guild_data["COMMANDS"] msg = f"Guild Free Roles: {free_roles}\n\nCommands:\n" for command_name, data in commands.items(): msg += f"\t{command_name}: {data['cost']}\n" msg += "\t\tRole Costs:\n" for role_id, cost in data.get("role_ids", {}).items(): role = guild.get_role(int(role_id)) if not role: continue msg += f"\t\t\t{role.name}: {cost}\n" msg += "\t\tUser Costs:\n" for user_id, cost in data.get("user_ids", {}).items(): user = guild.get_member(int(user_id)) if not user: continue msg += f"\t\t\t{user.name}: {cost}\n" msg = pagify(msg) for m in msg: await ctx.send(box(m, lang="python")) @commands.command(name="cost") @commands.guild_only() async def get_cost_command(self, ctx, command: str): """ Get cost of a command. """ if not self.arg_check(0, command): await ctx.send(warning(f"Command `{command}` not found!")) return cost = await self.get_cost(ctx, command=command) if cost == 0: await ctx.send(f"{command} is free for you!") return else: currency_name = await bank.get_currency_name(ctx.guild) await ctx.send(f"{command} costs `{cost}` {currency_name} for you.") @commands.Cog.listener() async def on_member_remove(self, member): await self.clean_data(member.guild) @commands.Cog.listener() async def on_guild_role_delete(self, role): await self.clean_data(role.guild) # Listens for poorerror @commands.Cog.listener() async def on_command_error(self, ctx, exception): if isinstance(exception, PoorError): cost = await self.get_cost(ctx) currency_name = await bank.get_currency_name(ctx.guild) balance = await bank.get_balance(ctx.author) message = await ctx.send( f"Sorry {ctx.author.name}, you do not have enough {currency_name} to use that command. (Cost: {cost}, Balance: {balance})", allowed_mentions=discord.AllowedMentions.all(), ) await asyncio.sleep(10) await message.delete() async def red_delete_data_for_user( self, *, requester: Literal["discord_deleted_user", "owner", "user", "user_strict"], user_id: int, ): pass