Implement /link & /unlink commands

Author: Lekuruu

app/bot.py

@@ -11,10 +11,12 @@ async def on_ready(self):
         app.session.logger.info(f'Logged in as {self.user}.')
         app.session.filters.populate()
         await self.load_cogs()
-        
+
     async def load_cogs(self):
         await self.load_extension("app.commands.kms")
+        await self.load_extension("app.commands.link")
         await self.load_extension("app.commands.bancho")
+        await self.tree.sync()
 
 def run():
     intents = discord.Intents.default()

app/cog.py

@@ -12,6 +12,7 @@
 class BaseCog(Cog):
     def __init__(self) -> None:
         self.bot = session.bot
+        self.guild = session.bot.guilds[0]
         self.redis = session.redis_async
         self.logger = session.logger
         self.events = session.events
@@ -25,4 +26,25 @@ async def run_async(func: Callable, *args):
         return await asyncio.get_event_loop().run_in_executor(None, func, *args)
 
     async def resolve_user(self, discord_id: int) -> DBUser | None:
-        return await self.run_async(users.fetch_by_discord_id, discord_id)
+        return await self.run_async(
+            users.fetch_by_discord_id,
+            discord_id
+        )
+
+    async def resolve_user_by_name(self, username: str) -> DBUser | None:
+        return await self.run_async(
+            users.fetch_by_name_extended,
+            username
+        )
+        
+    async def update_user(self, user_id: int, updates: dict) -> int:
+        return await self.run_async(
+            users.update,
+            user_id, updates
+        )
+
+    async def submit_event(self, name: str, *args) -> None:
+        return await self.run_async(
+            self.events.submit,
+            name, *args
+        )

app/commands/bancho.py

@@ -129,12 +129,6 @@ async def silence_user(self, user: DBUser, duration: float, reason: str) -> None
             user, duration, reason
         )
 
-    async def submit_event(self, name: str, *args) -> None:
-        return await self.run_async(
-            self.events.submit,
-            name, *args
-        )
-
     async def create_message(self, username: str, target: str, content: str) -> None:
         return await self.run_async(
             messages.create,

app/commands/link.py

@@ -0,0 +1,156 @@
+
+from discord.ext.commands import Cog, Bot, command
+from discord.ui import Modal, TextInput
+from discord import app_commands
+
+from app.common.database.objects import DBUser
+from app.common.cache import status
+from app.cog import BaseCog
+
+import discord
+import random
+import string
+import time
+
+class AccountLinking(BaseCog):
+    def __init__(self) -> None:
+        super().__init__()
+        self.member_role = discord.utils.get(self.guild.roles, name="Member")
+
+    @app_commands.command(name="link", description="Link your account to Titanic!")
+    @app_commands.describe(username="Your Titanic! username")
+    async def link_account(self, interaction: discord.Interaction, username: str) -> None:
+        if existing_user := await self.resolve_user(interaction.user.id):
+            return await interaction.response.send_message(
+                "Your account is already linked to Titanic! "
+                "Use /unlink to unlink your current account.",
+                ephemeral=True
+            )
+
+        if not (target_user := await self.resolve_user_by_name(username)):
+            return await interaction.response.send_message(
+                "No user found with that name.",
+                ephemeral=True
+            )
+
+        if target_user.discord_id:
+            return await interaction.response.send_message(
+                "This user is already linked to another Discord account.",
+                ephemeral=True
+            )
+
+        if not status.exists(target_user.id):
+            return await interaction.response.send_message(
+                "Please log into the game and try again!",
+                ephemeral=True
+            )
+
+        self.logger.info(f'[{interaction.user}] -> Starting linking process...')
+        
+        # Generate random 6-letter code which will be sent over DMs
+        code = ''.join(random.choices(string.ascii_lowercase, k=6))
+        await self.submit_event('link', target_user.id, code)
+
+        embed = discord.Embed(
+            title="🔗 Link Your Account",
+            description=(
+                f"You are linking the account: **{target_user.name}**\n"
+                "Click the button below to enter your code."
+            ),
+            color=discord.Color.blurple()
+        )
+        embed.set_footer(text="This message is only visible to you.")
+
+        await interaction.response.send_message(
+            view=LinkingView(code, target_user, self),
+            embed=embed,
+            ephemeral=True
+        )
+
+    @app_commands.command(name="unlink", description="Unlink your account from Titanic!")
+    async def unlink_account(self, interaction: discord.Interaction) -> None:
+        if not (linked_user := await self.resolve_user(interaction.user.id)):
+            return await interaction.response.send_message(
+                "Your account is not linked to Titanic!",
+                ephemeral=True
+            )
+
+        await self.update_user(
+            linked_user.id,
+            {"discord_id": None}
+        )
+
+        await interaction.response.send_message(
+            "You have successfully unlinked your account from Titanic!",
+            ephemeral=True
+        )
+
+    @Cog.listener()
+    async def on_member_join(self, member: discord.Member) -> None:
+        # Try to check if user already has their account linked to Titanic!
+        linked_user = await self.resolve_user(member.id)
+
+        if not linked_user:
+            return
+
+        # Re-add member role
+        await member.add_roles(self.member_role)
+
+class LinkingView(discord.ui.View):
+    def __init__(self, code: str, target_user: DBUser, cog: "AccountLinking"):
+        super().__init__(timeout=60*5)
+        self.target_user = target_user
+        self.code = code
+        self.cog = cog
+
+    @discord.ui.button(label="Enter Code", style=discord.ButtonStyle.primary)
+    async def enter_code(
+        self,
+        interaction: discord.Interaction,
+        button: discord.ui.Button
+    ) -> None:
+        modal = AccountLinkingModal(self.code, self.target_user, self.cog)
+        await interaction.response.send_modal(modal)
+
+class AccountLinkingModal(Modal):
+    def __init__(self, valid_code: str, target_user: DBUser, cog: "AccountLinking") -> None:
+        super().__init__(
+            title="Link your account",
+            timeout=60*5
+        )
+        self.code = TextInput(
+            label="Enter in-game code",
+            placeholder="abc123",
+            required=True,
+            max_length=6
+        )
+        self.add_item(self.code)
+        self.target_user = target_user
+        self.valid_code = valid_code
+        self.cog = cog
+
+    async def on_submit(self, interaction: discord.Interaction) -> None:
+        self.cog.logger.info(
+            f'[{interaction.user}] -> Entered code: "{self.code.value}"'
+        )
+
+        if self.code.value != self.valid_code:
+            return await interaction.response.send_message(
+                "Invalid code. Please try again.",
+                ephemeral=True
+            )
+
+        await self.cog.update_user(
+            self.target_user.id,
+            {"discord_id": interaction.user.id}
+        )
+        await interaction.response.send_message(
+            "Account linked successfully!",
+            ephemeral=True
+        )
+        self.cog.logger.info(
+            f'[{interaction.user}] -> Linked account: {self.target_user.name}'
+        )
+
+async def setup(bot: Bot):
+    await bot.add_cog(AccountLinking())