From d628bacef53cd881785ca82d629c745dcf08e199 Mon Sep 17 00:00:00 2001 From: Twentysix Date: Thu, 1 Dec 2016 02:30:51 +0100 Subject: [PATCH] [Core] Server based prefixes (#492) Added `[p]set serverprefix` `Bot.command_prefix` is now a callable that returns the prefixes set for the server. If none are set, it returns the global ones. --- cogs/alias.py | 28 +++++++------ cogs/audio.py | 4 +- cogs/customcom.py | 11 +++--- cogs/owner.py | 90 ++++++++++++++++++++++++++++-------------- cogs/utils/settings.py | 27 ++++++++++++- red.py | 29 +++++++++----- 6 files changed, 127 insertions(+), 62 deletions(-) diff --git a/cogs/alias.py b/cogs/alias.py index f048f4f3a..c39381889 100644 --- a/cogs/alias.py +++ b/cogs/alias.py @@ -3,8 +3,9 @@ from .utils.chat_formatting import * from .utils.dataIO import dataIO from .utils import checks from __main__ import user_allowed, send_cmd_help -import os from copy import deepcopy +import os +import discord class Alias: @@ -12,6 +13,7 @@ class Alias: self.bot = bot self.file_path = "data/alias/aliases.json" self.aliases = dataIO.load_json(self.file_path) + self.remove_old() @commands.group(pass_context=True, no_pm=True) async def alias(self, ctx): @@ -36,7 +38,7 @@ class Alias: await self.bot.say('I can\'t safely add an alias that starts with ' 'an existing command or alias. Sry <3') return - prefix = self.get_prefix(to_execute) + prefix = self.get_prefix(server, to_execute) if prefix is not None: to_execute = to_execute[len(prefix):] if server.id not in self.aliases: @@ -57,9 +59,10 @@ class Alias: server_aliases = self.aliases[server.id] if command in server_aliases: help_cmd = server_aliases[command].split(" ")[0] - new_content = self.bot.command_prefix[0] + new_content = self.bot.settings.get_prefixes(server)[0] new_content += "help " - new_content += help_cmd[len(self.get_prefix(help_cmd)):] + new_content += help_cmd[len(self.get_prefix(server, + help_cmd)):] message = ctx.message message.content = new_content await self.bot.process_commands(message) @@ -107,13 +110,13 @@ class Alias: else: await self.bot.say("There are no aliases on this server.") - async def check_aliases(self, message): + async def on_message(self, message): if len(message.content) < 2 or message.channel.is_private: return msg = message.content server = message.server - prefix = self.get_prefix(msg) + prefix = self.get_prefix(server, msg) if not prefix: return @@ -146,7 +149,8 @@ class Alias: if aliasname != self.first_word(aliasname): to_delete.append(aliasname) continue - prefix = self.get_prefix(alias) + server = discord.Object(id=sid) + prefix = self.get_prefix(server, alias) if prefix is not None: self.aliases[sid][aliasname] = alias[len(prefix):] for alias in to_delete: # Fixes caps and bad prefixes @@ -158,8 +162,9 @@ class Alias: def first_word(self, msg): return msg.split(" ")[0] - def get_prefix(self, msg): - for p in self.bot.command_prefix: + def get_prefix(self, server, msg): + prefixes = self.bot.settings.get_prefixes(server) + for p in prefixes: if msg.startswith(p): return p return None @@ -183,7 +188,4 @@ def check_file(): def setup(bot): check_folder() check_file() - n = Alias(bot) - n.remove_old() - bot.add_listener(n.check_aliases, "on_message") - bot.add_cog(n) + bot.add_cog(Alias(bot)) diff --git a/cogs/audio.py b/cogs/audio.py index eb2dd3879..e02c5d79c 100644 --- a/cogs/audio.py +++ b/cogs/audio.py @@ -745,8 +745,8 @@ class Audio: song = await self._guarantee_downloaded(server, url) except MaximumLength: log.warning("I can't play URL below because it is too long." - " Use {}audioset maxlength to change this.\n\n" - "{}".format(self.bot.command_prefix[0], url)) + " Use [p]audioset maxlength to change this.\n\n" + "{}".format(url)) raise local = False else: # Assume local diff --git a/cogs/customcom.py b/cogs/customcom.py index 157edb9a2..531b853a9 100644 --- a/cogs/customcom.py +++ b/cogs/customcom.py @@ -109,16 +109,15 @@ class CustomCommands: if len(message.content) < 2 or message.channel.is_private: return - msg = message.content server = message.server - prefix = self.get_prefix(msg) + prefix = self.get_prefix(message) if not prefix: return if server.id in self.c_commands and user_allowed(message): cmdlist = self.c_commands[server.id] - cmd = msg[len(prefix):] + cmd = message.content[len(prefix):] if cmd in cmdlist.keys(): cmd = cmdlist[cmd] cmd = self.format_cc(cmd, message) @@ -128,9 +127,9 @@ class CustomCommands: cmd = self.format_cc(cmd, message) await self.bot.send_message(message.channel, cmd) - def get_prefix(self, msg): - for p in self.bot.command_prefix: - if msg.startswith(p): + def get_prefix(self, message): + for p in self.bot.settings.get_prefixes(message.server): + if message.content.startswith(p): return p return False diff --git a/cogs/owner.py b/cogs/owner.py index 72a836009..116c55e39 100644 --- a/cogs/owner.py +++ b/cogs/owner.py @@ -1,7 +1,7 @@ import discord from discord.ext import commands from cogs.utils import checks -from __main__ import set_cog, send_cmd_help, settings +from __main__ import set_cog from .utils.dataIO import dataIO from .utils.chat_formatting import pagify, box @@ -13,7 +13,6 @@ import threading import datetime import glob import os -import time import aiohttp log = logging.getLogger("red.owner") @@ -226,7 +225,7 @@ class Owner: result = str(result) if not ctx.message.channel.is_private: - censor = (settings.email, settings.password) + censor = (self.bot.settings.email, self.bot.settings.password) r = "[EXPUNGED]" for w in censor: if w != "": @@ -258,7 +257,7 @@ class Owner: async def _set(self, ctx): """Changes Red's global settings.""" if ctx.invoked_subcommand is None: - await send_cmd_help(ctx) + await self.bot.send_cmd_help(ctx) return @_set.command(pass_context=True) @@ -268,7 +267,7 @@ class Owner: await self.bot.say("A set owner command is already pending.") return - if settings.owner != "id_here": + if self.bot.settings.owner != "id_here": await self.bot.say( "The owner is already set. Remember that setting the owner " "to someone else other than who hosts the bot has security " @@ -285,23 +284,52 @@ class Owner: @_set.command(pass_context=True) @checks.is_owner() async def prefix(self, ctx, *prefixes): - """Sets Red's prefixes + """Sets Red's global prefixes Accepts multiple prefixes separated by a space. Enclose in double quotes if a prefix contains spaces. Example: set prefix ! $ ? "two words" """ if prefixes == (): - await send_cmd_help(ctx) + await self.bot.send_cmd_help(ctx) return - self.bot.command_prefix = sorted(prefixes, reverse=True) - settings.prefixes = sorted(prefixes, reverse=True) - log.debug("Setting prefixes to:\n\t{}".format(settings.prefixes)) + self.bot.settings.prefixes = sorted(prefixes, reverse=True) + log.debug("Setting global prefixes to:\n\t{}" + "".format(self.bot.settings.prefixes)) - if len(prefixes) > 1: - await self.bot.say("Prefixes set") - else: - await self.bot.say("Prefix set") + p = "prefixes" if len(prefixes) > 1 else "prefix" + await self.bot.say("Global {} set".format(p)) + + @_set.command(pass_context=True, no_pm=True) + @checks.serverowner_or_permissions(administrator=True) + async def serverprefix(self, ctx, *prefixes): + """Sets Red's prefixes for this server + + Accepts multiple prefixes separated by a space. Enclose in double + quotes if a prefix contains spaces. + Example: set serverprefix ! $ ? "two words" + + Issuing this command with no parameters will reset the server + prefixes and the global ones will be used instead.""" + server = ctx.message.server + + if prefixes == (): + self.bot.settings.set_server_prefixes(server, []) + current_p = ", ".join(self.bot.settings.prefixes) + await self.bot.say("Server prefixes reset. Current prefixes: " + "`{}`".format(current_p)) + return + + prefixes = sorted(prefixes, reverse=True) + self.bot.settings.set_server_prefixes(server, prefixes) + log.debug("Setting server's {} prefixes to:\n\t{}" + "".format(server.id, self.bot.settings.prefixes)) + + p = "Prefixes" if len(prefixes) > 1 else "Prefix" + await self.bot.say("{} set for this server.\n" + "To go back to the global prefixes, do" + " `{}set serverprefix` " + "".format(p, prefixes[0])) @_set.command(pass_context=True) @checks.is_owner() @@ -310,16 +338,18 @@ class Owner: name = name.strip() if name != "": try: - await self.bot.edit_profile(settings.password, username=name) + await self.bot.edit_profile(self.bot.settings.password, + username=name) except: await self.bot.say("Failed to change name. Remember that you" " can only do it up to 2 times an hour." "Use nicknames if you need frequent " - "changes. {}set nickname".format(ctx.prefix)) + "changes. {}set nickname" + "".format(ctx.prefix)) else: await self.bot.say("Done.") else: - await send_cmd_help(ctx) + await self.bot.send_cmd_help(ctx) @_set.command(pass_context=True, no_pm=True) @checks.is_owner() @@ -391,7 +421,7 @@ class Owner: game=current_game) await self.bot.say("Status changed.") else: - await send_cmd_help(ctx) + await self.bot.send_cmd_help(ctx) @_set.command(pass_context=True) @checks.is_owner() @@ -412,7 +442,7 @@ class Owner: await self.bot.change_presence(game=game, status=current_status) log.debug('Owner has set streaming status and url to "{}" and {}'.format(stream_title, streamer)) elif streamer is not None: - await send_cmd_help(ctx) + await self.bot.send_cmd_help(ctx) return else: await self.bot.change_presence(game=None, status=current_status) @@ -426,7 +456,7 @@ class Owner: try: async with self.session.get(url) as r: data = await r.read() - await self.bot.edit_profile(settings.password, avatar=data) + await self.bot.edit_profile(self.bot.settings.password, avatar=data) await self.bot.say("Done.") log.debug("changed avatar") except Exception as e: @@ -442,9 +472,9 @@ class Owner: if len(token) < 50: await self.bot.say("Invalid token.") else: - settings.login_type = "token" - settings.email = token - settings.password = "" + self.bot.settings.login_type = "token" + self.bot.settings.email = token + self.bot.settings.password = "" await self.bot.say("Token set. Restart me.") log.debug("Token changed.") @@ -461,7 +491,7 @@ class Owner: With no subcommands returns the disabled commands list""" if ctx.invoked_subcommand is None: - await send_cmd_help(ctx) + await self.bot.send_cmd_help(ctx) if self.disabled_commands: msg = "Disabled commands:\n```xl\n" for cmd in self.disabled_commands: @@ -615,10 +645,11 @@ class Owner: @commands.command(pass_context=True) async def contact(self, ctx, *, message : str): """Sends message to the owner""" - if settings.owner == "id_here": + if self.bot.settings.owner == "id_here": await self.bot.say("I have no owner set.") return - owner = discord.utils.get(self.bot.get_all_members(), id=settings.owner) + owner = discord.utils.get(self.bot.get_all_members(), + id=self.bot.settings.owner) author = ctx.message.author if ctx.message.channel.is_private is False: server = ctx.message.server @@ -653,12 +684,13 @@ class Owner: py_version = "[{}.{}.{}]({})".format(*os.sys.version_info[:3], python_url) - owner = settings.owner if settings.owner != "id_here" else None + owner_set = self.bot.settings.owner != "id_here" + owner = self.bot.settings.owner if owner_set else None if owner: owner = discord.utils.get(self.bot.get_all_members(), id=owner) if not owner: try: - owner = await self.bot.get_user_info(settings.owner) + owner = await self.bot.get_user_info(self.bot.settings.owner) except: owner = None if not owner: @@ -749,7 +781,7 @@ class Owner: choice = input("> ") if choice == "yes": - settings.owner = author.id + self.bot.settings.owner = author.id print(author.name + " has been set as owner.") self.setowner_lock = False self.owner.hidden = True diff --git a/cogs/utils/settings.py b/cogs/utils/settings.py index aff380a05..1e90acf31 100644 --- a/cogs/utils/settings.py +++ b/cogs/utils/settings.py @@ -13,7 +13,9 @@ class Settings: self.default_settings = { "EMAIL": "EmailHere", "PASSWORD": "", "OWNER": "id_here", "PREFIXES": [], - "default": {"ADMIN_ROLE": "Transistor", "MOD_ROLE": "Process"}, + "default": {"ADMIN_ROLE": "Transistor", + "MOD_ROLE": "Process", + "PREFIXES": []}, "LOGIN_TYPE": "email"} if not dataIO.is_valid_json(self.path): self.bot_settings = self.default_settings @@ -46,7 +48,9 @@ class Settings: admin = self.bot_settings["ADMIN_ROLE"] del self.bot_settings["MOD_ROLE"] del self.bot_settings["ADMIN_ROLE"] - self.bot_settings["default"] = {"MOD_ROLE": mod, "ADMIN_ROLE": admin} + self.bot_settings["default"] = {"MOD_ROLE": mod, + "ADMIN_ROLE": admin, + "PREFIXES" : []} self.save_settings() @property @@ -171,6 +175,25 @@ class Settings: self.bot_settings[server.id]["MOD_ROLE"] = value self.save_settings() + def get_server_prefixes(self, server): + if server is None or server.id not in self.bot_settings: + return self.prefixes + return self.bot_settings[server.id].get("PREFIXES", []) + + def set_server_prefixes(self, server, prefixes): + if server is None: + return + assert isinstance(server, discord.Server) + if server.id not in self.bot_settings: + self.add_server(server.id) + self.bot_settings[server.id]["PREFIXES"] = prefixes + self.save_settings() + + def get_prefixes(self, server): + """Returns server's prefixes if set, otherwise global ones""" + p = self.get_server_prefixes(server) + return p if p else self.prefixes + def add_server(self, sid): self.bot_settings[sid] = self.bot_settings["default"].copy() self.save_settings() diff --git a/red.py b/red.py index c993e54be..3a40077e3 100644 --- a/red.py +++ b/red.py @@ -48,11 +48,23 @@ description = "Red - A multifunction Discord bot by Twentysix" class Bot(commands.Bot): def __init__(self, *args, **kwargs): + + def prefix_manager(bot, message): + """ + Returns prefixes of the message's server if set. + If none are set or if the message's server is None + it will return the global prefixes instead. + + Requires a Bot instance and a Message object to be + passed as arguments. + """ + return bot.settings.get_prefixes(message.server) + self.counter = Counter() self.uptime = datetime.datetime.now() self._message_modifiers = [] self.settings = Settings() - super().__init__(*args, **kwargs) + super().__init__(*args, command_prefix=prefix_manager, **kwargs) async def send_message(self, *args, **kwargs): if self._message_modifiers: @@ -173,8 +185,7 @@ class Formatter(commands.HelpFormatter): formatter = Formatter(show_check_failure=False) -bot = Bot(command_prefix=["_"], formatter=formatter, - description=description, pm_help=None) +bot = Bot(formatter=formatter, description=description, pm_help=None) send_cmd_help = bot.send_cmd_help # Backwards user_allowed = bot.user_allowed # compatibility @@ -200,8 +211,8 @@ async def on_ready(): print("{} users".format(users)) print("\n{}/{} active cogs with {} commands".format( len(bot.cogs), total_cogs, len(bot.commands))) - prefix_label = "Prefixes:" if len(bot.command_prefix) > 1 else "Prefix:" - print("{} {}\n".format(prefix_label, " ".join(bot.command_prefix))) + prefix_label = "Prefixes:" if len(settings.prefixes) > 1 else "Prefix:" + print("{} {}\n".format(prefix_label, " ".join(settings.prefixes))) if settings.login_type == "token": print("------") print("Use this url to bring your bot to a server:") @@ -477,18 +488,16 @@ def main(): check_configs() set_logger() owner_cog = load_cogs() - if settings.prefixes != []: - bot.command_prefix = settings.prefixes - else: + if settings.prefixes == []: print("No prefix set. Defaulting to !") - bot.command_prefix = ["!"] + settings.prefixes = ["!"] if settings.owner != "id_here": print("Use !set prefix to set it.") else: print("Once you're owner use !set prefix to set it.") if settings.owner == "id_here" and settings.login_type == "email": print("Owner has not been set yet. Do '{}set owner' in chat to set " - "yourself as owner.".format(bot.command_prefix[0])) + "yourself as owner.".format(settings.prefixes[0])) else: owner_cog.owner.hidden = True # Hides the set owner command from help print("-- Logging in.. --")