import re import discord from discord.ext.commands.context import Context import lavalink from discord.ext import commands from lavalink.models import DefaultPlayer from bot import ChristmasBot 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 # ensure there exists a client already if hasattr(self.client, "lavalink"): self.lavalink = self.client.lavalink else: self.client.lavalink = lavalink.Client(client.user.id) self.client.lavalink.add_node( "localhost", 2333, "youshallnotpass", "us", "default-node" ) 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(commands.Cog): def __init__(self, bot: ChristmasBot): self.bot = bot @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) 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 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, error): 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. 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 ("play",) 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: raise commands.CommandInvokeError("Not connected.") 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) @commands.command(name="play", aliases=["p", "connect", "join"]) @commands.guild_only() async def play(self, ctx: Context): """Searches and plays a song from a given query.""" # Get the player for this guild from cache. player: DefaultPlayer = self.bot.lavalink.player_manager.get(ctx.guild.id) query = "https://youtu.be/VcIt_AcOPjs" # Get the results for the query from Lavalink. results = await player.node.get_tracks(query) # Results could be None if Lavalink returns an invalid response (non-JSON/non-200 (OK)). # ALternatively, resullts['tracks'] could be an empty array if the query yielded no tracks. if not results or not results["tracks"]: return await EmbedGenerator.Title(ctx, "Nothing found!") # Theoretically songs will always be TRACK_LOADED track = results["tracks"][0] await EmbedGenerator.Message( ctx, "Track Enqueued", f'[{track["info"]["title"]}]({track["info"]["uri"]})' ) track = lavalink.models.AudioTrack(track, ctx.author.id, recommended=True) player.add(requester=ctx.author.id, track=track) if not player.is_playing: await player.play() @commands.command(name="disconnect", aliases=["dc", "stop"]) @commands.guild_only() async def disconnect(self, ctx: Context): """Disconnects the player from the voice channel and clears its queue.""" player: DefaultPlayer = self.bot.lavalink.player_manager.get(ctx.guild.id) if not player.is_connected: return await EmbedGenerator.Title(ctx, "Not connected.") 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. return await EmbedGenerator.Title(ctx, "You're not in my voicechannel!") # 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.") def setup(bot: ChristmasBot): bot.add_cog(Music(bot))