diff --git a/bot.py b/bot.py index 3b0a779..c03e31c 100644 --- a/bot.py +++ b/bot.py @@ -11,9 +11,13 @@ from discord.ext.commands.errors import ( ExtensionNotFound, NoEntryPointError, ) +import lavalink class ChristmasBot(commands.Bot): + lavalink: lavalink.Client + invite_link: str + def __init__(self, config: Dict[Any, Any]): intents = discord.Intents( voice_states=True, guild_messages=True, guilds=True, messages=True @@ -23,8 +27,6 @@ class ChristmasBot(commands.Bot): self.initial_cog_names: List[str] = self.config.get("cogs", []) self.colors: Dict[str, Color] = self.process_colours(config.get("colors", [])) - self.invite_link: str = "" - super().__init__( command_prefix=self.prefix_callable, description=self.config["info"].get("description", ""), diff --git a/cogs/music.py b/cogs/music.py index 855a346..f16309e 100644 --- a/cogs/music.py +++ b/cogs/music.py @@ -1,45 +1,91 @@ -import discord -from discord import channel -from discord.ext import commands +import re +import discord +from discord.ext.commands.context import Context import lavalink -import asyncio +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 MusicCog(commands.Cog, name="Music"): - def __init__(self, bot: commands.Bot): + +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 - self.stream = "https://azuracast.exobot.site/radio/8000/radio.opus" - if not hasattr(bot, "lavalink"): - bot.lavalink = lavalink.Client(bot.user.id) - bot.lavalink.add_node("de-1.rivalmc.net", 2333, "12345", "eu", "poggers") - bot.add_listener( - self.bot.lavalink.voice_update_handler, "on_socket_response" + @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"] ) - lavalink.add_event_hook(self.track_hook) - - bot.loop.create_task(self.async_init()) - - async def async_init(self): - await self.bot.wait_until_ready() - - channels = [] - - # We startup to fast #NOTPOGGERS - await asyncio.sleep(5) - for channel in channels: - guild = self.bot.get_guild(channel[0]) - player = self.bot.lavalink.player_manager.create( - channel[0], endpoint=str(guild.region) - ) - track = await player.node.get_tracks(self.stream) - if not player.is_playing: - await player.play(track["tracks"][0]) - await self.connect_to(channel[0], channel[1]) + self.bot.lavalink.add_event_hook(self.track_hook) def cog_unload(self): """Cog unload handler. This removes any event hooks that were registered.""" @@ -48,18 +94,42 @@ class MusicCog(commands.Cog, name="Music"): 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) ) - should_connect = ctx.command.name in ("connect",) + # 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: @@ -76,69 +146,74 @@ class MusicCog(commands.Cog, name="Music"): ) player.store("channel", ctx.channel.id) - await self.connect_to(ctx.guild.id, str(ctx.author.voice.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) - await self.connect_to(guild_id, None) + guild = self.bot.get_guild(guild_id) + await guild.voice_client.disconnect(force=True) - async def connect_to(self, guild_id: int, channel_id: str): - """Connects to the given voicechannel ID. A channel_id of `None` means disconnect.""" - ws = self.bot._connection._get_websocket(guild_id) - await ws.voice_state(str(guild_id), channel_id) + @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" - @commands.command(name="connect") - async def connect(self, ctx): - """Starts vibing.""" - player = self.bot.lavalink.player_manager.get(ctx.guild.id) - results = await player.node.get_tracks(self.stream) + # 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 ctx.send("Nothing found!") + return await EmbedGenerator.Title(ctx, "Nothing found!") - if results["loadType"] == "PLAYLIST_LOADED": - tracks = results["tracks"] + # Theoretically songs will always be TRACK_LOADED + track = results["tracks"][0] + await EmbedGenerator.Message( + ctx, "Track Enqueued", f'[{track["info"]["title"]}]({track["info"]["uri"]})' + ) - for track in tracks: - player.add(requester=ctx.author.id, track=track) - else: - track = results["tracks"][0] - track = lavalink.models.AudioTrack(track, ctx.author.id, recommended=True) - player.add(requester=ctx.author.id, track=track) + 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(aliases=["dc"]) - async def disconnect(self, ctx): + @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 = self.bot.lavalink.player_manager.get(ctx.guild.id) + player: DefaultPlayer = self.bot.lavalink.player_manager.get(ctx.guild.id) if not player.is_connected: - return await ctx.send("Not 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) ): - return await ctx.send("You're not in my voicechannel!") + # 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() - await self.connect_to(ctx.guild.id, None) - - @commands.command(name="now", aliases=["playing"]) - async def now_playing(self, ctx): - """Stop and disconnect the player and controller.""" - em = discord.Embed(color=self.bot.colors["embed"]) - # em.set_thumbnail(url=np["thumbnail"]) - em.add_field(name="Currently playing:", value="Some song") - await EmbedGenerator.SendWithFooter(ctx, em) + # Disconnect from the voice channel. + await ctx.voice_client.disconnect(force=True) + await EmbedGenerator.Title(ctx, "*⃣ | Disconnected.") def setup(bot: ChristmasBot): - bot.add_cog(MusicCog(bot)) + bot.add_cog(Music(bot))