import asyncio import datetime import re from typing import Optional import discord from discord.channel import TextChannel from discord.ext.commands.context import Context from discord.ext.commands.errors import CommandError import lavalink from discord import Embed from discord.channel import TextChannel from discord.ext import commands from discord.ext.commands.context import Context from lavalink.models import AudioTrack from lavalink.models import DefaultPlayer from bot import ChristmasBot from utils.classes import BaseCog from context import CustomContext from utils.database import AutoJoin from utils.database import Playlist from utils.exceptions import EmbeddedCommandException from utils.EmbedGenerator import EmbedGenerator url_rx = re.compile(r"https?://(?:www\.)?.+") class LavalinkVoiceClient(discord.VoiceClient): def __init__(self, client: discord.Client, channel: discord.abc.Connectable): self.client = client self.channel = channel self.lavalink = self.client.lavalink async def on_voice_server_update(self, data): # the data needs to be transformed before being handed down to # voice_update_handler lavalink_data = {"t": "VOICE_SERVER_UPDATE", "d": data} await self.lavalink.voice_update_handler(lavalink_data) async def on_voice_state_update(self, data): # the data needs to be transformed before being handed down to # voice_update_handler lavalink_data = {"t": "VOICE_STATE_UPDATE", "d": data} await self.lavalink.voice_update_handler(lavalink_data) async def connect(self, *, timeout: float, reconnect: bool) -> None: """ Connect the bot to the voice channel and create a player_manager if it doesn't exist yet. """ # ensure there is a player_manager when creating a new voice_client self.lavalink.player_manager.create(guild_id=self.channel.guild.id) await self.channel.guild.change_voice_state(channel=self.channel) async def disconnect(self, *, force: bool) -> None: """ Handles the disconnect. Cleans up running player and leaves the voice client. """ player = self.lavalink.player_manager.get(self.channel.guild.id) # no need to disconnect if we are not connected if not force and not player.is_connected: return # None means disconnect await self.channel.guild.change_voice_state(channel=None) # update the channel_id of the player to None # this must be done because the on_voice_state_update that # would set channel_id to None doesn't get dispatched after the # disconnect player.channel_id = None self.cleanup() class Music(BaseCog): @commands.Cog.listener() async def on_ready(self): if not hasattr( self.bot, "lavalink" ): # This ensures the client isn't overwritten during cog reloads. self.bot.lavalink = lavalink.Client(self.bot.user.id) ll = self.bot.config["lavalink"] self.bot.lavalink.add_node( ll["host"], ll["port"], ll["password"], ll["region"], ll["name"] ) self.bot.lavalink.add_event_hook(self.track_hook) await self.async_init() async def async_init(self): redis_result = await AutoJoin.get_channels(self.bot._redis_client) while len(self.bot.lavalink.node_manager.available_nodes) == 0: await asyncio.sleep(1) for guild_id, (voicechannel_id, textchannel_id) in redis_result.items(): player = self.bot.lavalink.player_manager.create(guild_id) player.store("channel", textchannel_id) voice_channel = await self.bot.fetch_channel(voicechannel_id) await voice_channel.connect(cls=LavalinkVoiceClient) if not player.is_playing: await self.fill_player_queue( player, self.bot.config["queue_buffer_size"] ) await player.play() textchannel = await self.bot.fetch_channel(textchannel_id) await textchannel.send("Automatically joined the voice channel") async def fill_player_queue(self, player: DefaultPlayer, buffer: Optional[int] = 1): queries = await Playlist.random(self.bot._redis_client, buffer) # Get the results for the query from Lavalink. for query in queries: result = await player.node.get_tracks(query) if not result or not result["tracks"]: continue track = lavalink.models.AudioTrack( result["tracks"][0], self.bot.user.id, recommended=False ) player.add(requester=self.bot.user.id, track=track) async def create_track_embed(self, track: AudioTrack) -> Embed: embed_color = self.bot.colors["embed"] embed = discord.Embed( title=f"Now playing...", colour=embed_color, ) embed.description = f"[{track.title}]({track.uri})" embed.set_thumbnail( url=f"https://i3.ytimg.com/vi/{track.identifier}/mqdefault.jpg" ) try: duration = str(datetime.timedelta(milliseconds=int(track.duration))) except OverflowError: duration = "🔴 LIVE" embed.add_field(name="Duration", value=duration) embed.add_field(name="Author", value=track.author) return embed def cog_unload(self): """Cog unload handler. This removes any event hooks that were registered.""" self.bot.lavalink._event_hooks.clear() async def cog_before_invoke(self, ctx): """Command before-invoke handler.""" guild_check = ctx.guild is not None # This is essentially the same as `@commands.guild_only()` # except it saves us repeating ourselves (and also a few lines). if not hasattr(self.bot, "lavalink"): await ctx.send("Still starting please wait a moment.") if guild_check: await self.ensure_voice(ctx) # Ensure that the bot and command author share a mutual voicechannel. return guild_check async def cog_command_error(self, ctx: CustomContext, error: CommandError): if isinstance(error, commands.CommandInvokeError): await ctx.send(error.original) # The above handles errors thrown in this cog and shows them to the user. # This shouldn't be a problem as the only errors thrown in this cog are from `ensure_voice` # which contain a reason string, such as "Join a voicechannel" etc. You can modify the above # if you want to do things differently. elif isinstance(error, EmbeddedCommandException): await error.send(ctx) async def ensure_voice(self, ctx): """This check ensures that the bot and command author are in the same voicechannel.""" player = self.bot.lavalink.player_manager.create( ctx.guild.id, endpoint=str(ctx.guild.region) ) # Create returns a player if one exists, otherwise creates. # This line is important because it ensures that a player always exists for a guild. # Most people might consider this a waste of resources for guilds that aren't playing, but this is # the easiest and simplest way of ensuring players are created. # These are commands that require the bot to join a voicechannel (i.e. initiating playback). # Commands such as volume/skip etc don't require the bot to be in a voicechannel so don't need listing here. should_connect = ctx.command.name in ("connect",) if not ctx.author.voice or not ctx.author.voice.channel: # Our cog_command_error handler catches this and sends it to the voicechannel. # Exceptions allow us to "short-circuit" command invocation via checks so the # execution state of the command goes no further. raise commands.CommandInvokeError("Join a voicechannel first.") if not player.is_connected: if not should_connect: bot_name = self.bot.config["info"]["name"] embed = await EmbedGenerator.Message( ctx, f"{bot_name} is not connected", "However you can start playing music using `/connect`", no_send=True, ) raise EmbeddedCommandException(embed) permissions = ctx.author.voice.channel.permissions_for(ctx.me) if ( not permissions.connect or not permissions.speak ): # Check user limit too? raise commands.CommandInvokeError( "I need the `CONNECT` and `SPEAK` permissions." ) player.store("channel", ctx.channel.id) await ctx.author.voice.channel.connect(cls=LavalinkVoiceClient) else: if int(player.channel_id) != ctx.author.voice.channel.id: raise commands.CommandInvokeError("You need to be in my voicechannel.") async def track_hook(self, event): if isinstance(event, lavalink.events.QueueEndEvent): # When this track_hook receives a "QueueEndEvent" from lavalink.py # it indicates that there are no tracks left in the player's queue. # To save on resources, we can tell the bot to disconnect from the voicechannel. guild_id = int(event.player.guild_id) guild = self.bot.get_guild(guild_id) await guild.voice_client.disconnect(force=True) elif isinstance(event, lavalink.events.TrackStartEvent): channel_id = int(event.player.fetch("channel")) channel: TextChannel = self.bot.get_channel(channel_id) embed = await self.create_track_embed(event.player.current) await channel.send(embed=embed) elif isinstance(event, lavalink.events.TrackEndEvent): await self.fill_player_queue(event.player, 1) @commands.command(name="connect", aliases=["p", "play", "join"]) async def play(self, ctx: CustomContext): """Start the radio""" # Get the player for this guild from cache. player = ctx.get_player() if player.is_connected: await ctx.send("Already connected") return await self.fill_player_queue(player, self.bot.config["queue_buffer_size"]+1) if not player.is_playing: await player.play() await EmbedGenerator.Title(ctx, "*⃣ | Connected.") return @commands.command(name="skip", aliases=["next"]) async def skip(self, ctx: CustomContext): """Skip the current song""" player = ctx.get_player() await player.skip() await ctx.send("Skipped current song") @commands.command(name="queue") async def queue(self, ctx: CustomContext): """Display the current radio queue""" player = ctx.get_player() embed_color = self.bot.colors["embed"] embed = Embed(title="Coming Up...", colour=embed_color) if len(player.queue) > 0: embed.description = "\n".join( f"{index}. [{track.title}]({track.uri})" for index, track in enumerate(player.queue, 1) ) else: embed.description = "We are still determining a playlist" await ctx.send(embed=embed) @commands.command(name="disconnect", aliases=["dc", "stop"]) async def disconnect(self, ctx: CustomContext): """Disconnects the radio from the channel""" player = ctx.get_player() if not ctx.author.voice or ( player.is_connected and ctx.author.voice.channel.id != int(player.channel_id) ): # Abuse prevention. Users not in voice channels, or not in the same voice channel as the bot # may not disconnect the bot. await EmbedGenerator.Title(ctx, "You're not in my voicechannel!") return # Clear the queue to ensure old tracks don't start playing # when someone else queues something. player.queue.clear() # Stop the current track so Lavalink consumes less resources. await player.stop() # Disconnect from the voice channel. await ctx.voice_client.disconnect(force=True) await EmbedGenerator.Title(ctx, "*⃣ | Disconnected.") @commands.command(name="now") async def now_playing(self, ctx: CustomContext): player: DefaultPlayer = self.bot.lavalink.player_manager.get(ctx.guild.id) track = player.current embed = await self.create_track_embed(track) await ctx.send(embed=embed) def setup(bot: ChristmasBot): bot.add_cog(Music(bot))