Merge pull request #27 from roblabla/switchroot

Add invite correlation system

Robocop.py

@@ -40,7 +40,8 @@ def get_prefix(bot, message):
 
 wanted_jsons = ["data/restrictions.json",
                 "data/robocronptab.json",
-                "data/userlog.json"]
+                "data/userlog.json",
+                "data/invites.json"]
 
 initial_extensions = ['cogs.common',
                       'cogs.admin',
@@ -59,7 +60,8 @@ initial_extensions = ['cogs.common',
                       'cogs.remind',
                       'cogs.robocronp',
                       'cogs.meme',
-                      'cogs.pin']
+                      'cogs.pin',
+                      'cogs.invites']
 
 bot = commands.Bot(command_prefix=get_prefix,
                    description=config.bot_description, pm_help=True)

cogs/invites.py

@@ -0,0 +1,43 @@
+from discord.ext import commands
+from discord.ext.commands import Cog
+from helpers.checks import check_if_collaborator
+import config
+import json
+
+class Invites(Cog):
+    def __init__(self, bot):
+        self.bot = bot
+
+    @commands.command()
+    @commands.guild_only()
+    @commands.check(check_if_collaborator)
+    async def invite(self, ctx):
+        welcome_channel = self.bot.get_channel(config.welcome_channel)
+        author = ctx.message.author
+        reason = f"Created by {str(author)} ({author.id})"
+        invite = await welcome_channel.create_invite(max_age = 0,
+                max_uses = 1, temporary = True, unique = True, reason = reason)
+
+        with open("data/invites.json", "r") as f:
+            invites = json.load(f)
+
+        invites[invite.id] = {
+            "uses": 0,
+            "url": invite.url,
+            "max_uses": 1,
+            "code": invite.code
+        }
+
+        with open("data/invites.json", "w") as f:
+            f.write(json.dumps(invites))
+
+        await ctx.message.add_reaction("🆗")
+        try:
+            await ctx.author.send(f"Created single-use invite {invite.url}")
+        except discord.errors.Forbidden:
+            await ctx.send(f"{ctx.author.mention} I could not send you the \
+                             invite. Send me a DM so I can reply to you.")
+
+
+def setup(bot):
+    bot.add_cog(Invites(bot))

cogs/logs.py

@@ -28,6 +28,54 @@ class Logs(Cog):
         # We use this a lot, might as well get it once
         escaped_name = self.bot.escape_message(member)
 
+        # Attempt to correlate the user joining with an invite
+        with open("data/invites.json", "r") as f:
+            invites = json.load(f)
+
+        real_invites = await member.guild.invites()
+
+        # Add unknown active invites. Can happen if invite was manually created
+        for invite in real_invites:
+            if invite.id not in invites:
+                invites[invite.id] = {
+                    "uses": 0,
+                    "url": invite.url,
+                    "max_uses": invite.max_uses,
+                    "code": invite.code
+                }
+
+        probable_invites_used = []
+        items_to_delete = []
+        # Look for invites whose usage increased since last lookup
+        for id, invite in invites.items():
+            real_invite = next((x for x in real_invites if x.id == id), None)
+
+            if real_invite is None:
+                # Invite does not exist anymore. Was either revoked manually
+                # or the final use was used up
+                probable_invites_used.append(invite)
+                items_to_delete.append(id)
+            elif invite["uses"] < real_invite.uses:
+                probable_invites_used.append(invite)
+                invite["uses"] = real_invite.uses
+
+        # Delete used up invites
+        for id in items_to_delete:
+            del invites[id]
+
+        # Save invites data.
+        with open("data/invites.json", "w") as f:
+            f.write(json.dumps(invites))
+
+        # Prepare the invite correlation message
+        if len(probable_invites_used) == 1:
+            invite_used = probable_invites_used[0]["url"]
+        elif len(probable_invites_used) == 0:
+            invite_used = "Unknown"
+        else:
+            invite_used = "One of: "
+            invite_used += ", ".join([x["code"] for x in probable_invites_used])
+
         # Check if user account is older than 15 minutes
         age = member.joined_at - member.created_at
         if age < config.min_age:
@@ -39,10 +87,12 @@ class Logs(Cog):
             except discord.errors.Forbidden:
                 sent = False
             await member.kick(reason="Too new")
+
             msg = f"🚨 **Account too new**: {member.mention} | "\
                   f"{escaped_name}\n"\
                   f"🗓 __Creation__: {member.created_at}\n"\
                   f"🕓 Account age: {age}\n"\
+                  f"✉  Joined with: {invite_used}\n"\
                   f"🏷 __User ID__: {member.id}"
             if not sent:
                 msg += "\nThe user has disabled direct messages,"\
@@ -53,6 +103,7 @@ class Logs(Cog):
               f"{escaped_name}\n"\
               f"🗓 __Creation__: {member.created_at}\n"\
               f"🕓 Account age: {age}\n"\
+              f"✉  Joined with: {invite_used}\n"\
               f"🏷 __User ID__: {member.id}"
 
         # Handles user restrictions

cogs/pin.py

@@ -1,7 +1,12 @@
 import config
+from discord.ext import commands
 from discord.ext.commands import Cog
 from discord.enums import MessageType
-
+from discord import Embed
+import aiohttp
+import gidgethub.aiohttp
+from helpers.checks import check_if_collaborator
+from helpers.checks import check_if_pin_channel
 
 class Pin(Cog):
     """
@@ -11,16 +16,83 @@ class Pin(Cog):
     def __init__(self, bot):
         self.bot = bot
 
+    def is_pinboard(self, msg):
+        return msg.author == self.bot.user and \
+            len(msg.embeds) > 0 and \
+            msg.embeds[0].title == "Pinboard"
+
+    async def get_pinboard(self, gh, channel):
+        # Find pinboard pin
+        pinboard_msg = None
+        for msg in reversed(await channel.pins()):
+            if self.is_pinboard(msg):
+                # Found pinboard, return content and gist id
+                id = msg.embeds[0].url.split("/")[-1]
+                data = await gh.getitem(f"/gists/{id}")
+                return (id, data["files"]["pinboard.md"]["content"])
+
+        # Create pinboard pin if it does not exist
+        data = await gh.post("/gists", data={
+            "files": {
+                "pinboard.md": {
+                    "content": "Old pins are available here:\n\n"
+                }
+            },
+            "description": f"Pinboard for SwitchRoot #{channel.name}",
+            "public": True
+        })
+
+        msg = await channel.send(embed=Embed(
+            title="Pinboard",
+            description="Old pins are moved to the pinboard to make space for \
+                         new ones. Check it out!",
+            url=data["html_url"]))
+        await msg.pin()
+
+        return (data["id"], data["files"]["pinboard.md"]["content"])
+
+    async def add_pin_to_pinboard(self, channel, data):
+        if config.github_oauth_token == "":
+            # Don't add to gist pinboard if we don't have an oauth token
+            return
+
+        async with aiohttp.ClientSession() as session:
+            gh = gidgethub.aiohttp.GitHubAPI(session, "RoboCop-NG",
+                    oauth_token=config.github_oauth_token)
+            (id, content) = await self.get_pinboard(gh, channel)
+            content += "- " + data + "\n"
+
+            await gh.patch(f"/gists/{id}", data={
+                "files": {
+                    "pinboard.md": {
+                        "content": content
+                    }
+                }
+            })
+
+    @commands.command()
+    @commands.guild_only()
+    @commands.check(check_if_collaborator)
+    @commands.check(check_if_pin_channel)
+    async def unpin(self, ctx, idx: int):
+        """Unpins a pinned message."""
+        if idx <= 50:
+            # Get message by pin idx
+            target_msg = (await ctx.message.channel.pins())[idx]
+        else:
+            # Get message by ID
+            target_msg = await ctx.message.channel.get_message(idx)
+        if self.is_pinboard(target_msg):
+            await ctx.send("Cannot unpin pinboard!")
+        else:
+            await target_msg.unpin()
+            await target_msg.remove_reaction("📌", self.bot.user)
+            await ctx.send(f"Unpinned {target_msg.jump_url}")
+            # TODO: Remove from pinboard?
+
     # Use raw_reaction to allow pinning old messages.
     @Cog.listener()
     async def on_raw_reaction_add(self, payload):
-        # TODO: handle more than 50 pinned message
-        # BODY: If there are more than 50 pinned messages,
-        # BODY: we should move the oldest pin to a pinboard
-        # BODY: channel to make room for the new pin.
-        # BODY: This is why we use the pin reaction to remember
-        # BODY: that a message is pinned.
-
         # Check that the user wants to pin this message
         if payload.emoji.name not in ["📌", "📍"]:
             return
@@ -45,17 +117,33 @@ class Pin(Cog):
                     if reaction.emoji == "📌":
                         if reaction.me:
                             return
-                        break
+                        else:
+                            break
 
-                # Wait for the automated "Pinned" message so we can delete it
+                # Add pin to pinboard, create one if none is found
-                waitable = self.bot.wait_for('message', check=check)
+                await self.add_pin_to_pinboard(target_chan, target_msg.jump_url)
 
-                # Pin the message
+                # Avoid staying "stuck" waiting for the pin message if message
-                await target_msg.pin()
+                # was already manually pinned
+                if not target_msg.pinned:
+                    # If we already have 50 pins, we should unpin the oldest.
+                    # We should avoid unpinning the pinboard.
+                    pins = await target_chan.pins()
+                    if len(pins) >= 50:
+                        for msg in reversed(pins):
+                            if not self.is_pinboard(msg):
+                                await msg.unpin()
+                                break
 
-                # Delete the automated Pinned message
+                    # Wait for the automated "Pinned" message so we can delete it
-                msg = await waitable
+                    waitable = self.bot.wait_for('message', check=check)
-                await msg.delete()
+
+                    # Pin the message
+                    await target_msg.pin()
+
+                    # Delete the automated Pinned message
+                    msg = await waitable
+                    await msg.delete()
 
                 # Add a Pin reaction so we remember that the message is pinned
                 await target_msg.add_reaction("📌")

config_template.py

@@ -46,6 +46,7 @@ staff_role_ids = [364647829248933888,  # Team role in ReSwitched
 # Various log channels used to log bot and guild's activity
 # You can use same channel for multiple log types
 # Spylog channel logs suspicious messages or messages by members under watch
+# Invites created with .invite will direct to the welcome channel.
 log_channel = 290958160414375946  # server-logs in ReSwitched
 botlog_channel = 529070282409771048  # bot-logs channel in ReSwitched
 modlog_channel = 542114169244221452  # mod-logs channel in ReSwitched
@@ -93,3 +94,6 @@ spy_channels = general_channels
 # Channels and roles where users can pin messages
 allowed_pin_channels = []
 allowed_pin_roles = []
+
+# Used for the pinboard. Leave empty if you don't wish for a gist pinboard.
+github_oauth_token = ""

helpers/checks.py

@@ -1,6 +1,5 @@
 import config
 
-
 def check_if_staff(ctx):
     if not ctx.guild:
         return False
@@ -20,3 +19,11 @@ def check_if_staff_or_ot(ctx):
     is_bot_cmds = (ctx.channel.name == "bot-cmds")
     is_staff = any(r.id in config.staff_role_ids for r in ctx.author.roles)
     return (is_ot or is_staff or is_bot_cmds)
+
+def check_if_collaborator(ctx):
+    return any(r.id in config.staff_role_ids + config.allowed_pin_roles for r in ctx.author.roles)
+
+
+def check_if_pin_channel(ctx):
+    return ctx.message.channel.id in config.allowed_pin_channels
+

requirements.txt

@@ -3,4 +3,6 @@ git+https://github.com/Rapptz/discord.py@rewrite
 asyncio
 python-dateutil
 humanize
-parsedatetime
+parsedatetime
+aiohttp
+gidgethub