diff --git a/costmanager/__init__.py b/costmanager/__init__.py new file mode 100644 index 0000000..6723274 --- /dev/null +++ b/costmanager/__init__.py @@ -0,0 +1,4 @@ +from .manager import CostManager + +def setup(bot): + bot.add_cog(CostManager(bot)) diff --git a/costmanager/manager.py b/costmanager/manager.py new file mode 100644 index 0000000..fab8b49 --- /dev/null +++ b/costmanager/manager.py @@ -0,0 +1,395 @@ +from redbot.core.utils.chat_formatting import * +from redbot.core import Config, checks, commands, bank +from redbot.core.bot import Red +import discord + +from collections import defaultdict + +class PoorError(commands.CommandError): + 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 + """ + 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["role_ids"].keys()) + found_roles = charged_roles & member_roles + cost = 0 + + # check user cost + if str(member.id) in command_data["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.") + + 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)) + + if 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_free_roles(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["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["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") + 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) + await ctx.send(f"Sorry {ctx.author.name}, you do not have enough {currency_name} to use that command. (Cost: {cost}, Balance: {balance})")