restrictions: implement

This means that we have achieved feature parity with robocop.

README.md

@@ -9,6 +9,8 @@ Based on https://gitlab.com/ao/dpybotbase
 
 ## TODO
 
+ALL FEATURES OF KURISU/ROBOCOP USED IN RESWITCHED ARE NOW SUPPORTED!
+
 - [x] .py configs
 - [x] membercount command
 - [x] Meme commands and pegaswitch (honestly the easiest part)
@@ -34,7 +36,7 @@ Based on https://gitlab.com/ao/dpybotbase
 - [x] Moderation: botnickname
 - [x] Moderation: nickname
 - [x] Moderation: clear/purge
-- [ ] Moderation: restrictions (people who leave with muted role will get muted role on join)
+- [x] Moderation: restrictions (people who leave with muted role will get muted role on join)
 - [x] Warns: warn
 - [x] Warns: listwarns-listwarnsid
 - [x] Warns: clearwarns-clearwarnsid

cogs/logs.py

@@ -46,6 +46,16 @@ 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, name=rst))
+            await member.add_roles(*roles)
+
         # Real hell zone.
         with open("data/warnsv2.json", "r") as f:
             warns = json.load(f)

cogs/mod.py

@@ -15,13 +15,34 @@ class ModCog:
     def check_if_target_is_staff(self, target):
         return any(r.id in config.staff_role_ids for r in target.roles)
 
+    async def add_restriction(self, member, rst):
+        # from kurisu source, credits go to ihaveamac
+        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)
+
+    async def remove_restriction(self, member, rst):
+        # from kurisu source, credits go to ihaveamac
+        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)
+
     @commands.guild_only()
     @commands.bot_has_permissions(kick_members=True)
     @commands.check(check_if_staff)
     @commands.command()
     async def mute(self, ctx, target: discord.Member, *, reason: str = ""):
         """Mutes a user, staff only."""
-        # TODO: keep a restriction list
         # so that muted people can't just leave and rejoin
         if self.check_if_target_is_staff(target):
             return await ctx.send("I can't mute this user as "
@@ -57,6 +78,7 @@ class ModCog:
         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 self.add_restriction(target, config.mute_role_name)
 
     @commands.guild_only()
     @commands.bot_has_permissions(kick_members=True)
@@ -80,6 +102,7 @@ class ModCog:
         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 self.remove_restriction(target, config.mute_role_name)
 
     @commands.guild_only()
     @commands.bot_has_permissions(kick_members=True)

config.py.template

@@ -43,3 +43,4 @@ community_channels = [526378423468425236]  # Channels requiring community role
 general_channels = [526372255052201995]  # Channels everyone can access
 
 mute_role = 526500080879140874  # Mute role in NotSwitched
+mute_role_name = "muted"  # Mute role name in NotSwitched

data/restrictions.json

@@ -0,0 +1 @@
+{}