From 1f63a4c50dab2c9dee9ff98480996346dcd2d173 Mon Sep 17 00:00:00 2001 From: Lemon4ksan Date: Sat, 15 Mar 2025 21:43:41 +0300 Subject: [PATCH] impr: Remove restricted access to the Vibe. --- MusicBot/cogs/utils/base_bot.py | 49 +++++++++++------- MusicBot/cogs/utils/voice_extension.py | 69 ++++++++++++------------- MusicBot/ui/menu.py | 71 ++++++++++++++++---------- 3 files changed, 106 insertions(+), 83 deletions(-) diff --git a/MusicBot/cogs/utils/base_bot.py b/MusicBot/cogs/utils/base_bot.py index 251e3b9..762872f 100644 --- a/MusicBot/cogs/utils/base_bot.py +++ b/MusicBot/cogs/utils/base_bot.py @@ -35,10 +35,10 @@ class BaseBot: Returns: (YMClient | None): Client or None. """ - logging.debug("[VC_EXT] Initializing Yandex Music client") + logging.debug("[BASE_BOT] Initializing Yandex Music client") if not (token := await self.get_ym_token(ctx)): - logging.debug("[VC_EXT] No token found") + logging.debug("[BASE_BOT] No token found") await self.send_response_message(ctx, "❌ Укажите токен через /account login.", delete_after=15, ephemeral=True) return None @@ -138,6 +138,31 @@ class BaseBot: logging.debug(f"[BASE_BOT] Failed to get message: {e}") raise + async def get_discord_user_by_id(self, ctx: ApplicationContext | Interaction | RawReactionActionEvent, user_id: int) -> discord.User | None: + if isinstance(ctx, ApplicationContext) and ctx.user: + logging.debug(f"[BASE_BOT] Getting user {user_id} from ApplicationContext") + return await ctx.bot.fetch_user(user_id) + elif isinstance(ctx, Interaction): + logging.debug(f"[BASE_BOT] Getting user {user_id} from Interaction") + return await ctx.client.fetch_user(user_id) + elif not self.bot: + raise ValueError("Bot instance is not available") + else: + logging.debug(f"[BASE_BOT] Getting user {user_id} from bot instance") + return await self.bot.fetch_user(user_id) + + async def get_viber_id_from_ctx(self, ctx: ApplicationContext | Interaction | RawReactionActionEvent) -> int | None: + if not ctx.guild_id: + logging.warning("[BASE_BOT] Guild not found") + return None + + guild = await self.db.get_guild(ctx.guild_id, projection={'current_viber_id': 1}) + + if guild['current_viber_id']: + return guild['current_viber_id'] + + return ctx.user_id if isinstance(ctx, discord.RawReactionActionEvent) else ctx.user.id if ctx.user else None + async def update_menu_views_dict( self, ctx: ApplicationContext | Interaction | RawReactionActionEvent, @@ -152,31 +177,17 @@ class BaseBot: guild (ExplicitGuild): Guild. disable (bool, optional): Disable menu. Defaults to False. """ - logging.debug(f"[VC_EXT] Updating menu views dict for guild {ctx.guild_id}") + logging.debug(f"[BASE_BOT] Updating menu views dict for guild {ctx.guild_id}") from MusicBot.ui import MenuView if not ctx.guild_id: - logging.warning("[VC_EXT] Guild not found") + logging.warning("[BASE_BOT] Guild not found") return if ctx.guild_id in self.menu_views: self.menu_views[ctx.guild_id].stop() self.menu_views[ctx.guild_id] = await MenuView(ctx).init(disable=disable) - - async def get_discord_user_by_id(self, ctx: ApplicationContext | Interaction | RawReactionActionEvent, user_id: int) -> discord.User | None: - - if isinstance(ctx, ApplicationContext) and ctx.user: - logging.debug(f"[BASE_BOT] Getting user {user_id} from ApplicationContext") - return await ctx.bot.fetch_user(user_id) - elif isinstance(ctx, Interaction): - logging.debug(f"[BASE_BOT] Getting user {user_id} from Interaction") - return await ctx.client.fetch_user(user_id) - elif not self.bot: - raise ValueError("Bot instance is not available") - else: - logging.debug(f"[BASE_BOT] Getting user {user_id} from bot instance") - return await self.bot.fetch_user(user_id) def get_current_event_loop(self, ctx: ApplicationContext | Interaction | RawReactionActionEvent) -> asyncio.AbstractEventLoop: """Get the current event loop. If the context is a RawReactionActionEvent, get the loop from the self.bot instance. @@ -200,4 +211,4 @@ class BaseBot: raise ValueError("Bot is not set.") return self.bot.loop else: - raise TypeError(f"Invalid context type: '{type(ctx).__name__}'.") \ No newline at end of file + raise TypeError(f"Invalid context type: '{type(ctx).__name__}'.") diff --git a/MusicBot/cogs/utils/voice_extension.py b/MusicBot/cogs/utils/voice_extension.py index 65b2a4f..95b1495 100644 --- a/MusicBot/cogs/utils/voice_extension.py +++ b/MusicBot/cogs/utils/voice_extension.py @@ -264,21 +264,19 @@ class VoiceExtension(BaseBot): """ logging.info(f"[VC_EXT] Updating vibe for guild {ctx.guild_id} with type '{vibe_type}' and id '{item_id}'") - uid = viber_id if viber_id else ctx.user_id if isinstance(ctx, discord.RawReactionActionEvent) else ctx.user.id if ctx.user else None + uid = viber_id if viber_id else await self.get_viber_id_from_ctx(ctx) if not uid or not ctx.guild_id: logging.warning("[VC_EXT] Guild ID or User ID not found in context") return False - user = await self.users_db.get_user(uid, projection={'vibe_settings': 1}) - guild = await self.db.get_guild(ctx.guild_id, projection={'vibing': 1, 'current_track': 1}) - if not (client := await self.init_ym_client(ctx)): return False if update_settings: logging.debug("[VIBE] Updating vibe settings") + user = await self.users_db.get_user(uid, projection={'vibe_settings': 1}) settings = user['vibe_settings'] await client.rotor_station_settings2( f"{vibe_type}:{item_id}", @@ -287,6 +285,8 @@ class VoiceExtension(BaseBot): language=settings['lang'] ) + guild = await self.db.get_guild(ctx.guild_id, projection={'vibing': 1, 'current_track': 1}) + if not guild['vibing']: try: feedback = await client.rotor_station_feedback_radio_started( @@ -371,7 +371,7 @@ class VoiceExtension(BaseBot): guild = await self.db.get_guild(ctx.guild_id, projection={'current_viber_id': 1, 'vibing': 1}) if guild['vibing'] and ctx.user.id != guild['current_viber_id']: logging.debug("[VIBE] Context user is not the current viber") - await ctx.respond("❌ Вы не можете взаимодействовать с чужой волной!", delete_after=15, ephemeral=True) + await ctx.respond("❌ Вы не можете изменять чужую волну!", delete_after=15, ephemeral=True) return False logging.debug("[VC_EXT] Voice requirements met") @@ -504,7 +504,8 @@ class VoiceExtension(BaseBot): menu_message: discord.Message | None = None, button_callback: bool = False ) -> str | None: - """Switch to the next track in the queue. Return track title on success. Performs all additional actions like updating menu and sending vibe feedback. + """Switch to the next track in the queue. Return track title on success. + Performs all additional actions like updating menu and sending vibe feedback. Doesn't change track if stopped. Stop playing if tracks list is empty. Args: @@ -519,14 +520,11 @@ class VoiceExtension(BaseBot): """ logging.debug("[VC_EXT] Switching to next track") - uid = ctx.user_id if isinstance(ctx, discord.RawReactionActionEvent) else ctx.user.id if ctx.user else None - - if not ctx.guild_id or not uid: + if not (uid := await self.get_viber_id_from_ctx(ctx)) or not ctx.guild_id: logging.warning("[VC_EXT] Guild ID or User ID not found in context inside 'next_track'") return None guild = await self.db.get_guild(ctx.guild_id, projection={'shuffle': 1, 'repeat': 1, 'is_stopped': 1, 'current_menu': 1, 'vibing': 1, 'current_track': 1}) - user = await self.users_db.get_user(uid) if guild['is_stopped'] and after: logging.debug("[VC_EXT] Playback is stopped, skipping after callback.") @@ -553,7 +551,12 @@ class VoiceExtension(BaseBot): next_track = await self.db.get_track(ctx.guild_id, 'next') if not next_track and guild['vibing']: + # NOTE: Real vibe gets next tracks after each skip. For smoother experience + # we get next tracks only after all the other tracks are finished + logging.debug("[VC_EXT] No next track found, generating new vibe") + + user = await self.users_db.get_user(uid) if not user['vibe_type'] or not user['vibe_id']: logging.warning("[VC_EXT] No vibe type or vibe id found in user data") return None @@ -701,9 +704,6 @@ class VoiceExtension(BaseBot): bool: Success status. """ logging.info(f"[VOICE] Performing '{vote_data['action']}' action for message {ctx.message_id}") - - if guild['current_viber_id']: - ctx.user_id = guild['current_viber_id'] if not ctx.guild_id: logging.warning("[VOICE] Guild not found") @@ -818,20 +818,11 @@ class VoiceExtension(BaseBot): """ logging.debug(f"[VC_EXT] Sending vibe feedback, type: {feedback_type}") - uid = ctx.user_id if isinstance(ctx, discord.RawReactionActionEvent) else ctx.user.id if ctx.user else None - - if not uid or not ctx.guild_id: + if not (uid := await self.get_viber_id_from_ctx(ctx)) or not ctx.guild_id: logging.warning("[VC_EXT] User id or guild id not found") return False - guild = await self.db.get_guild(ctx.guild_id, projection={'current_viber_id': 1}) - - if guild['current_viber_id']: - viber_id = guild['current_viber_id'] - else: - viber_id = uid - - user = await self.users_db.get_user(viber_id, projection={'vibe_batch_id': 1, 'vibe_type': 1, 'vibe_id': 1}) + user = await self.users_db.get_user(uid, projection={'vibe_batch_id': 1, 'vibe_type': 1, 'vibe_id': 1}) if not (client := await self.init_ym_client(ctx)): logging.info(f"[VC_EXT] Failed to init YM client for user {user['_id']}") @@ -871,12 +862,18 @@ class VoiceExtension(BaseBot): logging.warning(f"[VC_EXT] Timed out while downloading track '{track.title}'") raise - async def _delete_menu_message(self, ctx: ApplicationContext | Interaction | RawReactionActionEvent, current_menu: int, gid: int) -> Literal[True]: + async def _delete_menu_message( + self, + ctx: ApplicationContext | Interaction | RawReactionActionEvent, + current_menu: int, + gid: int + ) -> Literal[True]: """Delete current menu message and stop menu view. Return True on success. Args: ctx (ApplicationContext | Interaction | RawReactionActionEvent): Context. - guild (ExplicitGuild): Guild. + current_menu (int): Current menu message ID. + gid (int): Guild ID. Returns: Literal[True]: Always returns True. @@ -916,34 +913,32 @@ class VoiceExtension(BaseBot): Returns: (str | None): Song title or None. """ - gid = ctx.guild_id - uid = ctx.user_id if isinstance(ctx, discord.RawReactionActionEvent) else ctx.user.id if ctx.user else None - if not gid or not uid: + if not ctx.guild_id: logging.warning("Guild ID or User ID not found in context") return None - guild = await self.db.get_guild(gid, projection={'current_menu': 1, 'vibing': 1, 'current_track': 1}) + guild = await self.db.get_guild(ctx.guild_id, projection={'current_menu': 1, 'vibing': 1, 'current_track': 1}) if not (vc := await self.get_voice_client(ctx) if not vc else vc): return None try: if not guild['current_track'] or track.id != guild['current_track']['id']: - await self._download_track(gid, track) + await self._download_track(ctx.guild_id, track) except yandex_music.exceptions.TimedOutError: if not retry: return await self._play_track(ctx, track, vc=vc, menu_message=menu_message, button_callback=button_callback, retry=True) - else: - await self.send_response_message(ctx, f"😔 Не удалось загрузить трек. Попробуйте сбросить меню.", delete_after=15) - logging.error(f"[VC_EXT] Failed to download track '{track.title}'") + + await self.send_response_message(ctx, f"😔 Не удалось загрузить трек. Попробуйте сбросить меню.", delete_after=15) + logging.error(f"[VC_EXT] Failed to download track '{track.title}'") return None - async with aiofiles.open(f'music/{gid}.mp3', "rb") as f: + async with aiofiles.open(f'music/{ctx.guild_id}.mp3', "rb") as f: track_bytes = io.BytesIO(await f.read()) song = discord.FFmpegPCMAudio(track_bytes, pipe=True, options='-vn -b:a 64k -filter:a "volume=0.15"') - await self.db.set_current_track(gid, track) + await self.db.set_current_track(ctx.guild_id, track) if menu_message or guild['current_menu']: # Updating menu message before playing to prevent delay and avoid FFMPEG lags. @@ -966,7 +961,7 @@ class VoiceExtension(BaseBot): return None logging.info(f"[VC_EXT] Playing track '{track.title}'") - await self.db.update(gid, {'is_stopped': False}) + await self.db.update(ctx.guild_id, {'is_stopped': False}) if guild['vibing']: await self.send_vibe_feedback(ctx, 'trackStarted', track) diff --git a/MusicBot/ui/menu.py b/MusicBot/ui/menu.py index 194ddec..c2a90dc 100644 --- a/MusicBot/ui/menu.py +++ b/MusicBot/ui/menu.py @@ -29,7 +29,7 @@ class ToggleButton(Button, VoiceExtension): await interaction.respond("❌ Что-то пошло не так. Попробуйте снова.", delete_after=15, ephemeral=True) return - if not await self.voice_check(interaction, check_vibe_privilage=True): + if not await self.voice_check(interaction): return guild = await self.db.get_guild(gid) @@ -72,7 +72,8 @@ class PlayPauseButton(Button, VoiceExtension): async def callback(self, interaction: Interaction) -> None: logging.info('[MENU] Play/Pause button callback...') - if not await self.voice_check(interaction, check_vibe_privilage=True): + + if not await self.voice_check(interaction): return if not (gid := interaction.guild_id) or not interaction.user: @@ -114,11 +115,15 @@ class PlayPauseButton(Button, VoiceExtension): await interaction.respond("❌ Нет воспроизводимого трека.", delete_after=15, ephemeral=True) return + guild = await self.db.get_guild(interaction.guild_id, projection={'single_token_uid': 1}) + if vc.is_paused(): vc.resume() - embed.remove_footer() + if guild['single_token_uid'] and (user := await self.get_discord_user_by_id(interaction, guild['single_token_uid'])): + embed.set_footer(text=f"Используется токен {user.display_name}", icon_url=user.display_avatar.url) + else: + embed.remove_footer() else: - vc.pause() embed.set_footer(text='Приостановлено') await interaction.edit(embed=embed) @@ -139,7 +144,7 @@ class SwitchTrackButton(Button, VoiceExtension): logging.info(f'[MENU] {callback_type.capitalize()} track button callback') - if not await self.voice_check(interaction, check_vibe_privilage=True): + if not await self.voice_check(interaction): return tracks_type = callback_type + '_tracks' @@ -240,8 +245,7 @@ class LyricsButton(Button, VoiceExtension): if not await self.voice_check(interaction) or not interaction.guild_id or not interaction.user: return - client = await self.init_ym_client(interaction) - if not client: + if not (client := await self.init_ym_client(interaction)): return current_track = await self.db.get_track(interaction.guild_id, 'current') @@ -286,18 +290,18 @@ class MyVibeButton(Button, VoiceExtension): member = cast(Member, interaction.user) channel = cast(VoiceChannel, interaction.channel) track = await self.db.get_track(interaction.guild_id, 'current') - + if len(channel.members) > 2 and not member.guild_permissions.manage_channels: logging.info(f"Starting vote for starting vibe in guild {interaction.guild_id}") if track: response_message = f"{member.mention} хочет запустить волну по треку **{track['title']}**.\n\n Выполнить действие?" - _type = 'track' - _id = track['id'] + vibe_type = 'track' + vibe_id = track['id'] else: response_message = f"{member.mention} хочет запустить станцию **Моя Волна**.\n\n Выполнить действие?" - _type = 'user' - _id = 'onyourwave' + vibe_type = 'user' + vibe_id = 'onyourwave' message = cast(Interaction, await interaction.respond(response_message)) response = await message.original_response() @@ -313,7 +317,7 @@ class MyVibeButton(Button, VoiceExtension): 'negative_votes': list(), 'total_members': len(channel.members), 'action': 'vibe_station', - 'vote_content': [_type, _id, interaction.user.id] + 'vote_content': [vibe_type, vibe_id, interaction.user.id] } ) return @@ -337,8 +341,7 @@ class MyVibeButton(Button, VoiceExtension): logging.info('[MENU] Failed to start the vibe') await interaction.respond('❌ Не удалось запустить "Мою Волну". Возможно, у вас нет подписки на Яндекс Музыку.', ephemeral=True) - next_track = await self.db.get_track(interaction.guild_id, 'next') - if next_track: + if (next_track := await self.db.get_track(interaction.guild_id, 'next')): await self.play_track(interaction, next_track, button_callback=True) class MyVibeSelect(Select, VoiceExtension): @@ -539,8 +542,7 @@ class AddToPlaylistButton(Button, VoiceExtension): await interaction.respond('❌ Нет воспроизводимого трека.', delete_after=15, ephemeral=True) return - client = await self.init_ym_client(interaction) - if not client: + if not (client := await self.init_ym_client(interaction)): await interaction.respond('❌ Что-то пошло не так. Попробуйте позже.', delete_after=15, ephemeral=True) return @@ -582,7 +584,7 @@ class MenuView(View, VoiceExtension): self.play_pause_button = PlayPauseButton(style=ButtonStyle.primary, emoji='⏯', row=0) self.next_button = SwitchTrackButton(style=ButtonStyle.primary, emoji='⏭', row=0, custom_id='next') self.prev_button = SwitchTrackButton(style=ButtonStyle.primary, emoji='⏮', row=0, custom_id='previous') - + self.like_button = ReactionButton(style=ButtonStyle.secondary, emoji='❤️', row=1, custom_id='like') self.dislike_button = ReactionButton(style=ButtonStyle.secondary, emoji='💔', row=1, custom_id='dislike') self.lyrics_button = LyricsButton(style=ButtonStyle.secondary, emoji='📋', row=1) @@ -594,26 +596,33 @@ class MenuView(View, VoiceExtension): if not self.ctx.guild_id: return self - self.guild = await self.db.get_guild(self.ctx.guild_id, projection={'repeat': 1, 'shuffle': 1, 'current_track': 1, 'current_menu': 1, 'vibing': 1}) - + self.guild = await self.db.get_guild(self.ctx.guild_id, projection={ + 'repeat': 1, 'shuffle': 1, 'current_track': 1, 'current_menu': 1, 'vibing': 1, 'single_token_uid': 1 + }) + if self.guild['repeat']: self.repeat_button.style = ButtonStyle.success if self.guild['shuffle']: self.shuffle_button.style = ButtonStyle.success - + current_track = self.guild['current_track'] - likes = await self.get_liked_tracks(self.ctx) self.add_item(self.repeat_button) self.add_item(self.prev_button) self.add_item(self.play_pause_button) self.add_item(self.next_button) self.add_item(self.shuffle_button) - - if not isinstance(self.ctx, RawReactionActionEvent) and len(cast(VoiceChannel, self.ctx.channel).members) == 2: - if current_track and str(current_track['id']) in [str(like.id) for like in likes]: + + if not isinstance(self.ctx, RawReactionActionEvent) \ + and len(cast(VoiceChannel, self.ctx.channel).members) == 2 \ + and not self.guild['single_token_uid']: + + if current_track and str(current_track['id']) in [str(like.id) for like in await self.get_reacted_tracks(self.ctx, 'like')]: self.like_button.style = ButtonStyle.success + if current_track and str(current_track['id']) in [str(dislike.id) for dislike in await self.get_reacted_tracks(self.ctx, 'dislike')]: + self.dislike_button.style = ButtonStyle.success + if not current_track: self.lyrics_button.disabled = True self.like_button.disabled = True @@ -621,6 +630,11 @@ class MenuView(View, VoiceExtension): self.add_to_playlist_button.disabled = True elif not current_track['lyrics_available']: self.lyrics_button.disabled = True + + if self.guild['single_token_uid']: + self.like_button.disabled = True + self.dislike_button.disabled = True + self.add_to_playlist_button.disabled = True self.add_item(self.like_button) self.add_item(self.dislike_button) @@ -643,8 +657,11 @@ class MenuView(View, VoiceExtension): return if self.guild['current_menu']: - await self.stop_playing(self.ctx) - await self.db.update(self.ctx.guild_id, {'current_menu': None, 'previous_tracks': [], 'vibing': False}) + await self.db.update(self.ctx.guild_id, { + 'current_menu': None, 'repeat': False, 'shuffle': False, + 'previous_tracks': [], 'next_tracks': [], 'votes': {}, + 'vibing': False, 'current_viber_id': None + }) if (message := await self.get_menu_message(self.ctx, self.guild['current_menu'])): await message.delete()