diff --git a/cogs/logs.py b/cogs/logs.py index eacb794..fd80204 100644 --- a/cogs/logs.py +++ b/cogs/logs.py @@ -1,6 +1,7 @@ import discord import json import config +from helpers.restrictions import get_user_restrictions class Logs: @@ -43,15 +44,11 @@ class Logs: f"🕓 Account age: {age}\n"\ f"🏷 __User ID__: {member.id}" - # Taken from kurisu source. - # Blame ihaveamac, not me. - with open("data/restrictions.json", "r") as f: - rsts = json.load(f) - if str(member.id) in rsts: - roles = [] - for rst in rsts[str(member.id)]: - roles.append(discord.utils.get(member.guild.roles, id=rst)) - await member.add_roles(*roles) + # Handles user restrictions + # Basically, gives back muted role to users that leave with it. + rsts = get_user_restrictions(member.id) + roles = [discord.utils.get(member.guild.roles, id=rst) for rst in rsts] + await member.add_roles(*roles) # Real hell zone. with open("data/userlog.json", "r") as f: diff --git a/cogs/mod.py b/cogs/mod.py index e81b3d6..f6bf177 100644 --- a/cogs/mod.py +++ b/cogs/mod.py @@ -57,7 +57,7 @@ class Mod: log_channel = self.bot.get_channel(config.log_channel) await log_channel.send(chan_message) await ctx.send(f"{target.mention} can no longer speak.") - await add_restriction(target, config.mute_role) + add_restriction(target.id, config.mute_role) @commands.guild_only() @commands.check(check_if_staff) @@ -76,7 +76,7 @@ class Mod: log_channel = self.bot.get_channel(config.log_channel) await log_channel.send(chan_message) await ctx.send(f"{target.mention} can now speak again.") - await remove_restriction(target, config.mute_role) + remove_restriction(target.id, config.mute_role) @commands.guild_only() @commands.bot_has_permissions(kick_members=True) diff --git a/helpers/restrictions.py b/helpers/restrictions.py index 4e5af9c..3f889e5 100644 --- a/helpers/restrictions.py +++ b/helpers/restrictions.py @@ -1,25 +1,42 @@ import json -def add_restriction(self, member, rst): - # from kurisu source, credits go to ihaveamac +def get_restrictions(): with open("data/restrictions.json", "r") as f: - rsts = json.load(f) - if str(member.id) not in rsts: - rsts[str(member.id)] = [] - if rst not in rsts[str(member.id)]: - rsts[str(member.id)].append(rst) - with open("data/restrictions.json", "w") as f: - json.dump(rsts, f) + return json.load(f) -def remove_restriction(self, member, rst): - # from kurisu source, credits go to ihaveamac +def set_restrictions(contents): + with open("data/restrictions.json", "w") as f: + f.write(contents) + + +def get_user_restrictions(uid): + uid = str(uid) with open("data/restrictions.json", "r") as f: rsts = json.load(f) - if str(member.id) not in rsts: - rsts[str(member.id)] = [] - if rst in rsts[str(member.id)]: - rsts[str(member.id)].remove(rst) - with open("data/restrictions.json", "w") as f: - json.dump(rsts, f) + if uid in rsts: + return rsts[uid] + return [] + + +def add_restriction(uid, rst): + # mostly from kurisu source, credits go to ihaveamac + uid = str(uid) + rsts = get_restrictions() + if uid not in rsts: + rsts[uid] = [] + if rst not in rsts[uid]: + rsts[uid].append(rst) + set_restrictions(json.dumps(rsts)) + + +def remove_restriction(uid, rst): + # mostly from kurisu source, credits go to ihaveamac + uid = str(uid) + rsts = get_restrictions() + if uid not in rsts: + rsts[uid] = [] + if rst in rsts[uid]: + rsts[uid].remove(rst) + set_restrictions(json.dumps(rsts))