Add support for APPLE MUSIC!

code/bot.py

@@ -37,7 +37,7 @@ class MyBot(commands.Bot):
 
     async def on_ready(self):
         config.LOG.info(f"{bot.user} has connected to Discord.")
-        config.LOG.info("Startup complete. Sync slash commands by DMing the bot ***sync")
+        config.LOG.info(f"Startup complete. Sync slash commands by DMing the bot {bot.command_prefix}tree sync (guild id)")
 
 
 bot = MyBot()

code/cogs/play.py

@@ -2,18 +2,22 @@ import discord
 import datetime
 from discord import app_commands
 from discord.ext import commands
-from cogs.music import Music
 from lavalink import LoadType
 import re
-from cogs.music import LavalinkVoiceClient
 import requests
 
-from config import BOT_COLOR
+from cogs.music import Music, LavalinkVoiceClient
-from custom_source import CustomSource
+from config import BOT_COLOR, APPLE_MUSIC_KEY
+from custom_sources import SpotifySource, AppleSource
 
 
 url_rx = re.compile(r"https?://(?:www\.)?.+")
 
+apple_headers = {
+    "Authorization": f"Bearer {APPLE_MUSIC_KEY}",
+    "Origin": "https://apple.com",
+}
+
 
 class Play(commands.Cog):
     def __init__(self, bot):
@@ -26,7 +30,7 @@ class Play(commands.Cog):
         "Play a song from your favorite music provider"
         player = self.bot.lavalink.player_manager.get(interaction.guild.id)
 
-        # Notify users that YouTube and Apple Music links are not allowed
+        # Notify users that YouTube links are not allowed
 
         if "youtube.com" in query or "youtu.be" in query:
             embed = discord.Embed(
@@ -36,18 +40,107 @@ class Play(commands.Cog):
             )
             return await interaction.response.send_message(embed=embed, ephemeral=True)
 
-        if "music.apple.com" in query:
+        ###
-            embed = discord.Embed(
+        ### APPLE MUSIC links, perform API requests and load all tracks from the playlist/album/track
-                title="Apple Music Not Supported",
+        ###
-                description="Unfortunately, Apple Music does not allow bots to stream from their platform. Try sending a link for a different platform, or simply type the name of the song and I will automatically find it on a supported platform.",
-                color=BOT_COLOR,
-            )
-            return await interaction.response.send_message(embed=embed, ephemeral=True)
 
-        # If a Spotify link is found, act accordingly
+        if "music.apple.com" in query:
-        # We use a custom source for this (I tried the LavaSrc plugin, but Spotify
+            embed = discord.Embed(color=BOT_COLOR)
-        # links would just result in random shit being played whenever ytsearch was removed)
+
-        if "open.spotify.com" in query:
+            if "/playlist/" in query and "?i=" not in query:
+                playlist_id = query.split("/playlist/")[1].split("/")[1]
+                # Get all of the tracks in the playlist (limit at 250)
+                playlist_url = f"https://api.music.apple.com/v1/catalog/us/playlists/{playlist_id}/tracks?limit=100"
+                response = requests.get(playlist_url, headers=apple_headers)
+                if response.status_code == 200:
+                    playlist = response.json()
+                    # Get the general playlist info (name, artwork)
+                    playlist_info_url = f"https://api.music.apple.com/v1/catalog/us/playlists/{playlist_id}"
+                    playlist_info = requests.get(playlist_info_url, headers=apple_headers)
+                    playlist_info = playlist_info.json()
+                    artwork_url = playlist_info["data"][0]["attributes"]["artwork"]["url"].replace(
+                        "{w}x{h}", "300x300"
+                    )
+
+                    embed.title = "Playlist Queued"
+                    embed.description = f"**{playlist_info['data'][0]['attributes']['name']}**\n` {len(playlist['data'])} ` tracks\n\nQueued by: {interaction.user.mention}"
+                    embed.set_thumbnail(url=artwork_url)
+                    embed.set_footer(
+                        text=datetime.datetime.now(datetime.timezone.utc).strftime(
+                            "%Y-%m-%d %H:%M:%S"
+                        )
+                        + " UTC"
+                    )
+                    # Add small alert if the playlist is the max size
+                    if len(playlist["data"]) == 100:
+                        embed.description += "\n\n*This playlist is longer than the 100 song maximum. Only the first 100 songs will be queued.*"
+
+                    await interaction.response.send_message(embed=embed)
+
+                    tracks = await AppleSource.load_playlist(
+                        self, interaction.user, playlist
+                    )
+                    for track in tracks["tracks"]:
+                        player.add(requester=interaction.user, track=track)
+
+            # If there is an album, not a specific song within the album
+            if "/album/" in query and "?i=" not in query:
+                album_id = query.split("/album/")[1].split("/")[1]
+                album_url = f"https://api.music.apple.com/v1/catalog/us/albums/{album_id}"
+                response = requests.get(album_url, headers=apple_headers)
+                if response.status_code == 200:
+                    album = response.json()
+
+                    embed.title = "Album Queued"
+                    embed.description = f"**{album['data'][0]['attributes']['name']}** by **{album['data'][0]['attributes']['artistName']}**\n` {len(album['data'][0]['relationships']['tracks']['data'])} ` tracks\n\nQueued by: {interaction.user.mention}"
+                    embed.set_thumbnail(
+                        url=album["data"][0]["attributes"]["artwork"]["url"].replace(
+                            "{w}x{h}", "300x300"
+                        )
+                    )
+                    embed.set_footer(
+                        text=datetime.datetime.now(datetime.timezone.utc).strftime(
+                            "%Y-%m-%d %H:%M:%S"
+                        )
+                        + " UTC"
+                    )
+                    await interaction.response.send_message(embed=embed)
+
+                    tracks = await AppleSource.load_album(self, interaction.user, album)
+                    for track in tracks["tracks"]:
+                        player.add(requester=interaction.user, track=track)
+
+            # If there is a specific song
+            if "/album/" in query and "?i=" in query:
+                song_id = query.split("/album/")[1].split("?i=")[1]
+                song_url = f"https://api.music.apple.com/v1/catalog/us/songs/{song_id}"
+                response = requests.get(song_url, headers=apple_headers)
+                if response.status_code == 200:
+                    song = response.json()
+
+                    embed.title = "Song Queued"
+                    embed.description = f"**{song['data'][0]['attributes']['name']}** by **{song['data'][0]['attributes']['artistName']}**\n\nQueued by: {interaction.user.mention}"
+                    embed.set_thumbnail(
+                        url=song["data"][0]["attributes"]["artwork"]["url"].replace(
+                            "{w}x{h}", "300x300"
+                        )
+                    )
+                    embed.set_footer(
+                        text=datetime.datetime.now(datetime.timezone.utc).strftime(
+                            "%Y-%m-%d %H:%M:%S"
+                        )
+                        + " UTC"
+                    )
+                    await interaction.response.send_message(embed=embed)
+
+                    results = await AppleSource.load_item(self, interaction.user, song)
+                    player.add(requester=interaction.user, track=results.tracks[0])
+
+        ###
+        ### SPOTIFY links, perform API requests and load all tracks from the playlist/album/track
+        ###
+
+        elif "open.spotify.com" in query:
             embed = discord.Embed(color=BOT_COLOR)
 
             if "open.spotify.com/playlist" in query:
@@ -68,7 +161,7 @@ class Play(commands.Cog):
                     )
                     await interaction.response.send_message(embed=embed)
 
-                    tracks = await CustomSource.load_playlist(
+                    tracks = await SpotifySource.load_playlist(
                         self, interaction.user, playlist
                     )
                     for track in tracks["tracks"]:
@@ -92,7 +185,7 @@ class Play(commands.Cog):
                     )
                     await interaction.response.send_message(embed=embed)
 
-                    tracks = await CustomSource.load_album(
+                    tracks = await SpotifySource.load_album(
                         self, interaction.user, album
                     )
                     for track in tracks["tracks"]:
@@ -116,7 +209,7 @@ class Play(commands.Cog):
                     )
                     await interaction.response.send_message(embed=embed)
 
-                    results = await CustomSource.load_item(
+                    results = await SpotifySource.load_item(
                         self, interaction.user, track
                     )
                     player.add(requester=interaction.user, track=results.tracks[0])
@@ -128,7 +221,9 @@ class Play(commands.Cog):
                     embed=embed, ephemeral=True
                 )
 
-        # For anything else, do default searches and attempt to find track and play
+        ###
+        ### For anything else, use default Lavalink providers to search the query
+        ###
 
         else:
             if not url_rx.match(query):

code/config.py

@@ -31,6 +31,7 @@ FEEDBACK_CHANNEL_ID = None
 BUG_CHANNEL_ID = None
 SPOTIFY_CLIENT_ID = None
 SPOTIFY_CLIENT_SECRET = None
+APPLE_MUSIC_KEY = None
 OPENAI_API_KEY = None
 LAVALINK_HOST = None
 LAVALINK_PORT = None
@@ -70,6 +71,10 @@ def load_config():
             "SPOTIFY_CLIENT_SECRET": "",
         }
 
+        config["APPLE_MUSIC"] = {
+            "APPLE_MUSIC_KEY": "",
+        }
+
         config["OPENAI"] = {
             "OPENAI_API_KEY": "",
         }
@@ -96,7 +101,7 @@ Validate all of the options in the config.ini file.
 
 
 def validate_config(file_contents):
-    global TOKEN, BOT_COLOR, BOT_INVITE_LINK, FEEDBACK_CHANNEL_ID, BUG_CHANNEL_ID, SPOTIFY_CLIENT_ID, SPOTIFY_CLIENT_SECRET, OPENAI_API_KEY, LAVALINK_HOST, LAVALINK_PORT, LAVALINK_PASSWORD
+    global TOKEN, BOT_COLOR, BOT_INVITE_LINK, FEEDBACK_CHANNEL_ID, BUG_CHANNEL_ID, SPOTIFY_CLIENT_ID, SPOTIFY_CLIENT_SECRET, APPLE_MUSIC_KEY, OPENAI_API_KEY, LAVALINK_HOST, LAVALINK_PORT, LAVALINK_PASSWORD
     config = configparser.ConfigParser()
     config.read_string(file_contents)
 
@@ -106,7 +111,7 @@ def validate_config(file_contents):
     errors = 0
 
     # Make sure all sections are present
-    if ["BOT_INFO", "SPOTIFY", "OPENAI", "LAVALINK"] != config.sections():
+    if ["BOT_INFO", "SPOTIFY", "APPLE_MUSIC", "OPENAI", "LAVALINK"] != config.sections():
         sys.exit(
             LOG.critical(
                 "Missing sections in config.ini file. Delete the file and re-run the bot to generate a blank config.ini file."
@@ -127,6 +132,13 @@ def validate_config(file_contents):
             )
         )
 
+    if ["apple_music_key"] != config.options("APPLE_MUSIC"):
+        sys.exit(
+            LOG.critical(
+                "Missing options in APPLE_MUSIC section of config.ini file. Delete the file and re-run the bot to generate a blank config.ini file."
+            )
+        )
+
     if ["openai_api_key"] != config.options("OPENAI"):
         sys.exit(
             LOG.critical(
@@ -175,6 +187,7 @@ def validate_config(file_contents):
     TOKEN = config["BOT_INFO"]["TOKEN"]
     SPOTIFY_CLIENT_ID = config["SPOTIFY"]["SPOTIFY_CLIENT_ID"]
     SPOTIFY_CLIENT_SECRET = config["SPOTIFY"]["SPOTIFY_CLIENT_SECRET"]
+    APPLE_MUSIC_KEY = config["APPLE_MUSIC"]["APPLE_MUSIC_KEY"]
     OPENAI_API_KEY = config["OPENAI"]["OPENAI_API_KEY"]
     LAVALINK_HOST = config["LAVALINK"]["HOST"]
     LAVALINK_PORT = config["LAVALINK"]["PORT"]
@@ -194,7 +207,7 @@ Validate all of the environment variables.
 
 
 def validate_env_vars():
-    global TOKEN, BOT_COLOR, BOT_INVITE_LINK, FEEDBACK_CHANNEL_ID, BUG_CHANNEL_ID, SPOTIFY_CLIENT_ID, SPOTIFY_CLIENT_SECRET, OPENAI_API_KEY, LAVALINK_HOST, LAVALINK_PORT, LAVALINK_PASSWORD
+    global TOKEN, BOT_COLOR, BOT_INVITE_LINK, FEEDBACK_CHANNEL_ID, BUG_CHANNEL_ID, SPOTIFY_CLIENT_ID, SPOTIFY_CLIENT_SECRET, APPLE_MUSIC_KEY, OPENAI_API_KEY, LAVALINK_HOST, LAVALINK_PORT, LAVALINK_PASSWORD
 
     hex_pattern_one = "^#([A-Fa-f0-9]{6}|[A-Fa-f0-9]{3})$"
     hex_pattern_two = "^([A-Fa-f0-9]{6}|[A-Fa-f0-9]{3})$"
@@ -202,7 +215,7 @@ def validate_env_vars():
     errors = 0
 
     # Make sure all required variables are present in the environment
-    required_vars = ["TOKEN", "BOT_COLOR", "BOT_INVITE_LINK", "SPOTIFY_CLIENT_ID", "SPOTIFY_CLIENT_SECRET", "OPENAI_API_KEY", "LAVALINK_HOST", "LAVALINK_PORT", "LAVALINK_PASSWORD"]
+    required_vars = ["TOKEN", "BOT_COLOR", "BOT_INVITE_LINK", "SPOTIFY_CLIENT_ID", "SPOTIFY_CLIENT_SECRET", "APPLE_MUSIC_KEY", "OPENAI_API_KEY", "LAVALINK_HOST", "LAVALINK_PORT", "LAVALINK_PASSWORD"]
 
     for var in required_vars:
         if var not in os.environ:
@@ -254,6 +267,7 @@ def validate_env_vars():
     TOKEN = os.environ["TOKEN"]
     SPOTIFY_CLIENT_ID = os.environ["SPOTIFY_CLIENT_ID"]
     SPOTIFY_CLIENT_SECRET = os.environ["SPOTIFY_CLIENT_SECRET"]
+    APPLE_MUSIC_KEY = os.environ["APPLE_MUSIC_KEY"]
     OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
     LAVALINK_HOST = os.environ["LAVALINK_HOST"]
     LAVALINK_PORT = os.environ["LAVALINK_PORT"]

code/custom_sources.py

@@ -122,3 +122,82 @@ class SpotifySource(Source):
             )
 
         return LoadResult(LoadType.PLAYLIST, tracks, playlist_info=PlaylistInfo.none())
+
+
+"""
+Custom Source for Apple Music links
+"""
+
+
+class AppleSource(Source):
+    def __init__(self):
+        super().__init__(name="custom")
+
+    async def load_item(self, user, metadata):
+        track = CustomAudioTrack(
+            {  # Create an instance of our CustomAudioTrack.
+                "identifier": metadata["data"][0]["id"],
+                "isSeekable": True,
+                "author": metadata["data"][0]["attributes"]["artistName"],
+                "length": metadata["data"][0]["attributes"]["durationInMillis"],
+                "isStream": False,
+                "title": metadata["data"][0]["attributes"]["name"],
+                "uri": metadata["data"][0]["attributes"]["url"],
+                "duration": metadata["data"][0]["attributes"]["durationInMillis"],
+                "artworkUrl": metadata["data"][0]["attributes"]["artwork"]["url"].replace(
+                    "{w}x{h}", "300x300"
+                ),
+            },
+            requester=user,
+        )
+        return LoadResult(LoadType.TRACK, [track], playlist_info=PlaylistInfo.none())
+
+    async def load_album(self, user, metadata):
+        tracks = []
+        for track in metadata["data"][0]["relationships"]["tracks"][
+            "data"
+        ]:  # Loop through each track in the album.
+            tracks.append(
+                CustomAudioTrack(
+                    {  # Create an instance of our CustomAudioTrack.
+                        "identifier": track["id"],
+                        "isSeekable": True,
+                        "author": track["attributes"]["artistName"],
+                        "length": track["attributes"]["durationInMillis"],
+                        "isStream": False,
+                        "title": track["attributes"]["name"],
+                        "uri": track["attributes"]["url"],
+                        "duration": track["attributes"]["durationInMillis"],
+                        "artworkUrl": track["attributes"]["artwork"]["url"].replace(
+                            "{w}x{h}", "300x300"
+                        ),
+                    },
+                    requster=user,
+                )
+            )
+
+        return LoadResult(LoadType.PLAYLIST, tracks, playlist_info=PlaylistInfo.none())
+
+    async def load_playlist(self, user, metadata):
+        tracks = []
+        for track in metadata["data"]:  # Loop through each track in the playlist.
+            tracks.append(
+                CustomAudioTrack(
+                    {  # Create an instance of our CustomAudioTrack.
+                        "identifier": track["id"],
+                        "isSeekable": True,
+                        "author": track["attributes"]["artistName"],
+                        "length": track["attributes"]["durationInMillis"],
+                        "isStream": False,
+                        "title": track["attributes"]["name"],
+                        "uri": track["attributes"]["url"],
+                        "duration": track["attributes"]["durationInMillis"],
+                        "artworkUrl": track["attributes"]["artwork"]["url"].replace(
+                            "{w}x{h}", "300x300"
+                        ),
+                    },
+                    requster=user,
+                )
+            )
+
+        return LoadResult(LoadType.PLAYLIST, tracks, playlist_info=PlaylistInfo.none())