From 97f534eb1709652a30614acfe9ed8cde6302c242 Mon Sep 17 00:00:00 2001
From: TSRBerry <20988865+TSRBerry@users.noreply.github.com>
Date: Sat, 9 Mar 2024 21:41:42 +0100
Subject: [PATCH] Automatically update configured vanity URLs (#91)

---
 robocop_ng/cogs/vanity_url.py | 34 ++++++++++++++++++++++++++++++++++
 robocop_ng/config_template.py |  4 ++++
 2 files changed, 38 insertions(+)
 create mode 100644 robocop_ng/cogs/vanity_url.py

diff --git a/robocop_ng/cogs/vanity_url.py b/robocop_ng/cogs/vanity_url.py
new file mode 100644
index 0000000..caf30b7
--- /dev/null
+++ b/robocop_ng/cogs/vanity_url.py
@@ -0,0 +1,34 @@
+from discord import Guild
+from discord.ext import tasks
+from discord.ext.commands import Cog
+
+
+class VanityUrl(Cog):
+    def __init__(self, bot):
+        self.bot = bot
+        self.vanity_codes: dict[int, str] = self.bot.config.vanity_codes
+        self.check_changed_vanity_codes.start()
+
+    def cog_unload(self):
+        self.check_changed_vanity_codes.cancel()
+
+    async def update_vanity_code(self, guild: Guild, code: str):
+        if "VANITY_URL" in guild.features and guild.vanity_url_code != code:
+            await guild.edit(
+                reason="Configured vanity code was different", vanity_code=code
+            )
+
+    @Cog.listener()
+    async def on_guild_update(self, before: Guild, after: Guild):
+        if after.id in self.vanity_codes:
+            await self.update_vanity_code(after, self.vanity_codes[after.id])
+
+    @tasks.loop(hours=12)
+    async def check_changed_vanity_codes(self):
+        await self.bot.wait_until_ready()
+        for guild, vanity_code in self.vanity_codes.items():
+            await self.update_vanity_code(self.bot.get_guild(guild), vanity_code)
+
+
+async def setup(bot):
+    await bot.add_cog(VanityUrl(bot))
diff --git a/robocop_ng/config_template.py b/robocop_ng/config_template.py
index ab492e6..3d9cf2a 100644
--- a/robocop_ng/config_template.py
+++ b/robocop_ng/config_template.py
@@ -38,6 +38,7 @@ initial_cogs = [
     "cogs.links",
     "cogs.remind",
     "cogs.robocronp",
+    "cogs.vanity_url",
     "cogs.meme",
     "cogs.invites",
     "cogs.yubicootp",
@@ -64,6 +65,9 @@ min_age = datetime.timedelta(minutes=15)
 # The bot will only work in these guilds
 guild_whitelist = [269333940928512010]  # ReSwitched discord
 
+# Custom invite URL codes
+vanity_codes = {269333940928512010: "reswitched"}
+
 # Named roles to be used with .approve and .revoke
 # Example: .approve User hacker
 named_roles = {