From 9e5c2f5777b812d733963607eb47f6ab19ba6767 Mon Sep 17 00:00:00 2001 From: Interitio Date: Thu, 12 Jun 2025 23:10:14 +1000 Subject: [PATCH] feat(profiles): Add profile link UI. --- src/brand.py | 6 + src/core/cog.py | 6 + src/modules/profiles/cog.py | 21 ++ src/modules/profiles/ui/__init__.py | 1 + src/modules/profiles/ui/twitchlink.py | 329 ++++++++++++++++++++++++++ 5 files changed, 363 insertions(+) create mode 100644 src/brand.py create mode 100644 src/modules/profiles/ui/__init__.py create mode 100644 src/modules/profiles/ui/twitchlink.py diff --git a/src/brand.py b/src/brand.py new file mode 100644 index 0000000..b471430 --- /dev/null +++ b/src/brand.py @@ -0,0 +1,6 @@ +import discord + + +# Theme +MAIN_COLOUR = discord.Colour.from_str('#11EA11') +ACCENT_COLOUR = discord.Colour.from_str('#EA11EA') diff --git a/src/core/cog.py b/src/core/cog.py index b5eab52..6934762 100644 --- a/src/core/cog.py +++ b/src/core/cog.py @@ -11,6 +11,7 @@ from meta.app import shardname, appname from meta.logger import log_wrap from utils.lib import utc_now +from datamodels import DataModel from .data import CoreData logger = logging.getLogger(__name__) @@ -29,7 +30,9 @@ class CoreCog(LionCog): def __init__(self, bot: LionBot): self.bot = bot self.data = CoreData() + self.datamodel = DataModel() bot.db.load_registry(self.data) + bot.db.load_registry(self.datamodel) self.app_config: Optional[CoreData.AppConfig] = None self.bot_config: Optional[CoreData.BotConfig] = None @@ -43,6 +46,9 @@ class CoreCog(LionCog): self.app_config = await self.data.AppConfig.fetch_or_create(appname) self.bot_config = await self.data.BotConfig.fetch_or_create(appname) + await self.data.init() + await self.datamodel.init() + # Load the app command cache await self.reload_appcmd_cache() diff --git a/src/modules/profiles/cog.py b/src/modules/profiles/cog.py index 43c21b5..1c631af 100644 --- a/src/modules/profiles/cog.py +++ b/src/modules/profiles/cog.py @@ -21,6 +21,8 @@ from .data import ProfileData from .profile import UserProfile from .community import Community +from .ui import TwitchLinkStatic, TwitchLinkFlow + class ProfileCog(LionCog): def __init__(self, bot: LionBot): @@ -34,6 +36,8 @@ class ProfileCog(LionCog): async def cog_load(self): await self.data.init() + self.bot.add_view(TwitchLinkStatic(timeout=None)) + async def cog_check(self, ctx): return True @@ -197,6 +201,16 @@ class ProfileCog(LionCog): community = await Community.create_from_twitch(self.bot, user) return community + # ----- Admin Commands ----- + @cmds.hybrid_command( + name='linkoffer', + description="Send a message with a permanent button for profile linking" + ) + @appcmds.default_permissions(manage_guild=True) + async def linkoffer_cmd(self, ctx: LionContext): + view = TwitchLinkStatic(timeout=None) + await ctx.channel.send(embed=view.embed, view=view) + # ----- Profile Commands ----- @cmds.hybrid_group( name='profiles', @@ -217,6 +231,13 @@ class ProfileCog(LionCog): description="Link a twitch account to your current profile." ) async def profiles_link_twitch_cmd(self, ctx: LionContext): + if not ctx.interaction: + return + flowui = TwitchLinkFlow(self.bot, ctx.author, callerid=ctx.author.id) + await flowui.run(ctx.interaction) + await flowui.wait() + + async def old_profiles_link_twitch_cmd(self, ctx: LionContext): if not ctx.interaction: return diff --git a/src/modules/profiles/ui/__init__.py b/src/modules/profiles/ui/__init__.py new file mode 100644 index 0000000..359a69f --- /dev/null +++ b/src/modules/profiles/ui/__init__.py @@ -0,0 +1 @@ +from .twitchlink import TwitchLinkStatic, TwitchLinkFlow diff --git a/src/modules/profiles/ui/twitchlink.py b/src/modules/profiles/ui/twitchlink.py new file mode 100644 index 0000000..0138acc --- /dev/null +++ b/src/modules/profiles/ui/twitchlink.py @@ -0,0 +1,329 @@ +""" +UI Views for Twitch linkage. + +- Persistent view with interaction-button to enter link flow. +- We don't store the view, but we listen to interaction button id. +- Command to enter link flow. + +For link flow, send ephemeral embed with instructions and what to expect, with link button below. +After auth is granted through OAuth flow (or if not granted, e.g. on timeout or failure) +edit the embed to reflect auth situation. + +If migration occurred, add the migration text as a field to the embed. +""" +import asyncio +from datetime import timedelta +from enum import IntEnum +from typing import Optional + +import aiohttp +import discord +from discord.ui.button import button, Button +from discord.enums import ButtonStyle +from twitchAPI.helper import first + +from meta import LionBot +from meta.errors import SafeCancellation +from meta.logger import log_wrap + +from utils.ui import MessageUI +from utils.lib import MessageArgs, utc_now +from utils.ui.leo import LeoUI + +from modules.profiles.profile import UserProfile + +import brand + +from .. import logger + +class TwitchLinkStatic(LeoUI): + """ + Static UI whose only job is to display a persistent button + to ask people to connect their twitch account. + """ + def __init__(self, **kwargs): + super().__init__(**kwargs) + self._embed: Optional[discord.Embed] = None + + async def interaction_check(self, interaction: discord.Interaction): + return True + + @property + def embed(self) -> discord.Embed: + """ + This is the persistent message people will see with the button that starts the Oauth flow. + + Not sure what this should actually say, or whether it should be customisable via command. + + :TODO-MARKER: + """ + embed = discord.Embed( + title="Link your Twitch account!", + description=( + "To participate in the Dreamspace Adventure Game :TM:, " + "please start by pressing the button below to begin the login flow with Twitch!" + ), + colour=brand.ACCENT_COLOUR, + ) + return embed + + @embed.setter + def embed(self, value): + self._embed = value + + @button(label="Link", custom_id="BTN-LINK-TWITCH", style=ButtonStyle.green, emoji='🔗') + async def button_linker(self, interaction: discord.Interaction, btn: Button): + # Here we just reply to the interaction with the AuthFlow UI + # TODO + flowui = TwitchLinkFlow(interaction.client, interaction.user, callerid=interaction.user.id) + await flowui.run(interaction) + await flowui.wait() + + +class FlowState(IntEnum): + SETUP = -1 + WAITING = 0 + + CANCELLED = 1 + TIMEOUT = 2 + ERRORED = 3 + + WORKING = 9 + DONE = 10 + + +class TwitchLinkFlow(MessageUI): + def __init__(self, bot: LionBot, caller: discord.User | discord.Member, *args, **kwargs): + kwargs.setdefault('callerid', caller.id) + super().__init__(*args, **kwargs) + self.bot = bot + + self._auth_task = None + self._stage: FlowState = FlowState.SETUP + self.flow = None + self.authrow = None + self.user = caller + + self._info = None + self._migration_details = None + + # ----- UI API ----- + async def run(self, interaction: discord.Interaction, **kwargs): + await interaction.response.defer(ephemeral=True, thinking=True) + await self._start_flow() + await self.draw(interaction, **kwargs) + if self._stage is FlowState.ERRORED: + # This can happen if starting the flow failed + await self.close() + + async def _start_flow(self): + try: + self.flow = await self.bot.get_cog('TwitchAuthCog').start_auth() + except aiohttp.ClientError: + self._stage = FlowState.ERRORED + self._info = ( + "Could not establish a connection to the authentication server! " + "Please try again later~" + ) + logger.exception("Unexpected exception while starting authentication flow!", exc_info=True) + else: + self._stage = FlowState.WAITING + self._auth_task = asyncio.create_task(self._auth_flow()) + + async def _auth_flow(self): + """ + Run the flow and wait for a timeout, cancellation, or callback. + Update the message accordingly. + """ + assert self.flow is not None + try: + # TODO: Cancel this in cleanup + authrow = await asyncio.wait_for(self.flow.run(), timeout=180) + except asyncio.TimeoutError: + self._stage = FlowState.TIMEOUT + # Link Timed Out! + self._info = ( + "We didn't receive a response for three minutes so we closed the uplink " + "to keep your account safe! If you still want to connect, please try again!" + ) + await self.refresh() + await self.close() + except asyncio.CancelledError: + # Presumably the user exited or the bot is shutting down. + # Not safe to edit the message, but try and cleanup + await self.close() + except SafeCancellation as e: + logger.info("User or server cancelled authentication flow: ", exc_info=True) + # Uplink Cancelled! + self._info = ( + f"We couldn't complete the uplink!\nReason:*{e.msg}*" + ) + self._stage = FlowState.CANCELLED + await self.refresh() + await self.close() + else: + self._stage = FlowState.WORKING + self._info = ( + "Authentication complete! Connecting your Dreamspace account ...." + ) + self.authrow = authrow + await self.refresh() + await self._link_twitch(str(authrow.userid)) + await self.refresh() + await self.close() + + async def cleanup(self): + await super().cleanup() + if self._auth_task and not self._auth_task.cancelled(): + self._auth_task.cancel() + + async def _link_twitch(self, twitch_id: str): + """ + Link the caller's profile to the given twitch_id. + + Performs migration if needed. + """ + try: + twitch_user = await first(self.bot.twitch.get_users(user_ids=[twitch_id])) + except Exception: + logger.exception( + f"Looking up user {self.authrow} from Twitch authentication flow raised an error." + ) + self._stage = FlowState.ERRORED + self._info = "Failed to look up your user details from Twitch! Please try again later." + return + + if twitch_user is None: + logger.error( + f"User {self.authrow} obtained from Twitch authentication does not exist." + ) + self._stage = FlowState.ERRORED + self._info = "Authentication failed! Please try again later." + return + + profiles = self.bot.get_cog('ProfileCog') + userid = self.user.id + + caller_profile = await UserProfile.fetch_from_discordid(self.bot, userid) + twitch_profile = await UserProfile.fetch_from_twitchid(self.bot, twitch_id) + + succ_info = ( + f"Successfully established uplink to your Twitch account **{twitch_user.display_name}** " + "and transferred dreamspace data! Happy adventuring, and watch out for the grue~" + ) + # ::TODO-MARKER:: + + if twitch_profile is None: + if caller_profile is None: + # Neither profile exists + profile = await UserProfile.create_from_discord(self.bot, self.user) + await profile.attach_twitch(twitch_id) + + self._stage = FlowState.DONE + self._info = succ_info + else: + await caller_profile.attach_twitch(twitch_id) + + self._stage = FlowState.DONE + self._info = succ_info + else: + if caller_profile is None: + await twitch_profile.attach_discord(self.user.id) + + self._stage = FlowState.DONE + self._info = succ_info + elif twitch_profile.profileid == caller_profile.profileid: + self._stage = FlowState.CANCELLED + self._info = ( + f"The Twitch account **{twitch_user.display_name}** is already linked to your profile!" + ) + else: + # In this case we have conflicting profiles we need to migrate + try: + results = await profiles.migrate_profile(twitch_profile, caller_profile) + except Exception: + self._stage = FlowState.ERRORED + self._info = ( + "An issue was encountered while merging your account profiles! " + "The migration was rolled back, and not data has been lost.\n" + "The developer has been notified, please try again later!" + ) + logger.exception(f"Failed to migrate profiles {twitch_profile=} to {caller_profile=}") + else: + self._stage = FlowState.DONE + self._info = succ_info + self._migration_details = '\n'.join(results) + logger.info( + f"Migrated {twitch_profile=} to {caller_profile}. Info: {self._migration_details}" + ) + + # ----- UI Flow ----- + async def make_message(self) -> MessageArgs: + if self._stage is FlowState.SETUP: + raise ValueError("Making message before flow initialisation!") + assert self.flow is not None + + if self._stage is FlowState.WAITING: + # Message should be the initial request page + dur = discord.utils.format_dt(utc_now() + timedelta(seconds=179), style='R') + + title = "Press the button to login!" + desc = ( + "We have generated a custom secure link for you to connect your Twitch profile! " + "Press the button below and accept the connection in your browser, " + "and we will begin the transfer!\n" + f"(Note: The link expires {dur})" + ) + colour = brand.ACCENT_COLOUR + elif self._stage is FlowState.CANCELLED: + # Show cancellation message + # Show 'you can close this' + title = "Uplink Cancelled!" + desc = self._info + colour = discord.Colour.brand_red() + elif self._stage is FlowState.TIMEOUT: + title = "Link Timed Out" + desc = self._info + colour = discord.Colour.brand_red() + elif self._stage is FlowState.ERRORED: + title = "Something went wrong!" + desc = self._info + colour = discord.Colour.brand_red() + elif self._stage is FlowState.WORKING: + # We've received the auth, we are now doing migration + title = "Establishing Connection" + desc = self._info + colour = brand.ACCENT_COLOUR + elif self._stage is FlowState.DONE: + title = "Success!" + desc = self._info + colour = discord.Colour.brand_green() + else: + raise ValueError(f"Invalid stage value {self._stage}") + + embed = discord.Embed(title=title, description=desc, colour=colour, timestamp=utc_now()) + if self._migration_details: + embed.add_field( + name="Profile migration details", + value=self._migration_details + ) + return MessageArgs(embed=embed) + + async def refresh_layout(self): + # If we haven't received the auth callback yet, make the flow link button + if self.flow is None: + raise ValueError("Refreshing before flow initialisation!") + + if self._stage <= FlowState.WAITING: + flow_link = self.flow.auth.return_auth_url() + button = Button( + style=ButtonStyle.link, + url=flow_link, + label="Login With Twitch" + ) + self.set_layout((button,)) + else: + self.set_layout(()) + + async def reload(self): + pass