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,33 @@
+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 {author.name}#{author.discriminator} ({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(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,48 @@ 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: " + ", ".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 +81,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 +97,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

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

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
+