645 lines
22 KiB
Python
645 lines
22 KiB
Python
"""MIT License
|
|
|
|
Copyright (c) 2023 - present Vocard Development
|
|
|
|
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
of this software and associated documentation files (the "Software"), to deal
|
|
in the Software without restriction, including without limitation the rights
|
|
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
copies of the Software, and to permit persons to whom the Software is
|
|
furnished to do so, subject to the following conditions:
|
|
|
|
The above copyright notice and this permission notice shall be included in all
|
|
copies or substantial portions of the Software.
|
|
|
|
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
SOFTWARE.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
import re
|
|
import aiohttp
|
|
import logging
|
|
|
|
from discord import Client, Member
|
|
from discord.ext.commands import Bot
|
|
from typing import Dict, Optional, TYPE_CHECKING, Union, List
|
|
from urllib.parse import quote
|
|
|
|
from . import (
|
|
__version__,
|
|
spotify,
|
|
)
|
|
|
|
from .enums import SearchType, NodeAlgorithm
|
|
from .exceptions import (
|
|
InvalidSpotifyClientAuthorization,
|
|
NodeConnectionFailure,
|
|
NodeCreationError,
|
|
NodeException,
|
|
NodeNotAvailable,
|
|
NoNodesAvailable,
|
|
TrackLoadError
|
|
)
|
|
from .objects import Playlist, Track
|
|
from .utils import ExponentialBackoff, NodeStats, NodeInfo, Ping
|
|
from .enums import RequestMethod
|
|
|
|
if TYPE_CHECKING:
|
|
from .player import Player
|
|
|
|
SPOTIFY_URL_REGEX = re.compile(
|
|
r"https?://open.spotify.com/(?P<type>album|playlist|track|artist)/(?P<id>[a-zA-Z0-9]+)"
|
|
)
|
|
|
|
DISCORD_MP3_URL_REGEX = re.compile(
|
|
r"https?://cdn.discordapp.com/attachments/(?P<channel_id>[0-9]+)/"
|
|
r"(?P<message_id>[0-9]+)/(?P<file>[a-zA-Z0-9_.]+)+"
|
|
)
|
|
|
|
URL_REGEX = re.compile(
|
|
r"https?://(?:www\.)?.+"
|
|
)
|
|
|
|
NODE_VERSION = "v4"
|
|
|
|
class Node:
|
|
"""The base class for a node.
|
|
This node object represents a Lavalink node.
|
|
To enable Spotify searching, pass in a proper Spotify Client ID and Spotify Client Secret
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
pool,
|
|
bot: Bot,
|
|
host: str,
|
|
port: int,
|
|
password: str,
|
|
identifier: str,
|
|
secure: bool = False,
|
|
heartbeat: int = 30,
|
|
session: Optional[aiohttp.ClientSession] = None,
|
|
spotify_client_id: Optional[str] = None,
|
|
spotify_client_secret: Optional[str] = None,
|
|
resume_key: Optional[str] = None,
|
|
logger: Optional[logging.Logger] = None
|
|
):
|
|
self._bot: Bot = bot
|
|
self._host: str = host
|
|
self._port: int = port
|
|
self._pool: NodePool = pool
|
|
self._password: str = password
|
|
self._identifier: str = identifier
|
|
self._heartbeat: int = heartbeat
|
|
self._secure: bool = secure
|
|
self._logger: Optional[logging.Logger] = logger
|
|
|
|
self._websocket_uri: str = f"{'wss' if self._secure else 'ws'}://{self._host}:{self._port}/" + NODE_VERSION + "/websocket"
|
|
self._rest_uri: str = f"{'https' if self._secure else 'http'}://{self._host}:{self._port}"
|
|
|
|
self._session: aiohttp.ClientSession = session or aiohttp.ClientSession()
|
|
self._websocket: aiohttp.ClientWebSocketResponse = None
|
|
self._task: asyncio.Task = None
|
|
|
|
self.resume_key: str = resume_key or str(os.urandom(8).hex())
|
|
self._session_id: str = None
|
|
self._available: bool = None
|
|
|
|
self._headers: Dict[str, str] = {
|
|
"Authorization": self._password,
|
|
"User-Id": str(bot.user.id),
|
|
"Client-Name": f"Voicelink/{__version__}",
|
|
'Resume-Key': self.resume_key
|
|
}
|
|
|
|
self._players: Dict[int, Player] = {}
|
|
self._info: Optional[NodeInfo] = None
|
|
|
|
self._spotify_client_id: Optional[str] = spotify_client_id
|
|
self._spotify_client_secret: Optional[str] = spotify_client_secret
|
|
self._spotify_client: Optional[spotify.Client] = None
|
|
|
|
self._bot.add_listener(self._update_handler, "on_socket_response")
|
|
|
|
def __repr__(self):
|
|
return (
|
|
f"<Voicelink.node ws_uri={self._websocket_uri} rest_uri={self._rest_uri} "
|
|
f"player_count={len(self._players)}>"
|
|
)
|
|
|
|
def get_player(self, guild_id: int) -> Optional[Player]:
|
|
"""Takes a guild ID as a parameter. Returns a voicelink Player object."""
|
|
return self._players.get(guild_id, None)
|
|
|
|
@property
|
|
def spotify_client(self) -> Optional[spotify.Client]:
|
|
if not self._spotify_client:
|
|
self._spotify_client = spotify.Client(
|
|
self._spotify_client_id, self._spotify_client_secret
|
|
)
|
|
|
|
return self._spotify_client
|
|
|
|
@property
|
|
def is_connected(self) -> bool:
|
|
""""Property which returns whether this node is connected or not"""
|
|
return self._websocket is not None and not self._websocket.closed
|
|
|
|
|
|
@property
|
|
def stats(self) -> NodeStats:
|
|
"""Property which returns the node stats."""
|
|
return self._stats
|
|
|
|
@property
|
|
def players(self) -> Dict[int, Player]:
|
|
"""Property which returns a dict containing the guild ID and the player object."""
|
|
return self._players
|
|
|
|
@property
|
|
def bot(self) -> Bot:
|
|
"""Property which returns the discord.py client linked to this node"""
|
|
return self._bot
|
|
|
|
@property
|
|
def player_count(self) -> int:
|
|
"""Property which returns how many players are connected to this node"""
|
|
return len(self.players)
|
|
|
|
@property
|
|
def pool(self) -> NodePool:
|
|
"""Property which returns the pool this node is apart of"""
|
|
return self._pool
|
|
|
|
@property
|
|
def latency(self) -> float:
|
|
"""Property which returns the latency of the node"""
|
|
return Ping(self._host, port=self._port).get_ping()
|
|
|
|
async def _update_handler(self, data: dict) -> None:
|
|
await self._bot.wait_until_ready()
|
|
|
|
if not data:
|
|
return
|
|
|
|
if data["t"] == "VOICE_SERVER_UPDATE":
|
|
guild_id = int(data["d"]["guild_id"])
|
|
try:
|
|
player = self._players[guild_id]
|
|
await player.on_voice_server_update(data["d"])
|
|
except KeyError:
|
|
return
|
|
|
|
elif data["t"] == "VOICE_STATE_UPDATE":
|
|
if int(data["d"]["user_id"]) != self._bot.user.id:
|
|
return
|
|
|
|
guild_id = int(data["d"]["guild_id"])
|
|
try:
|
|
player = self._players[guild_id]
|
|
await player.on_voice_state_update(data["d"])
|
|
except KeyError:
|
|
return
|
|
|
|
async def _listen(self) -> None:
|
|
backoff = ExponentialBackoff(base=7)
|
|
|
|
while True:
|
|
try:
|
|
msg = await self._websocket.receive()
|
|
except:
|
|
break
|
|
if msg.type == aiohttp.WSMsgType.CLOSED:
|
|
self._available = False
|
|
|
|
retry = backoff.delay()
|
|
self._logger.info(f"Trying to reconnect node [{self._identifier}] with {round(retry)}s")
|
|
await asyncio.sleep(retry)
|
|
if not self.is_connected:
|
|
try:
|
|
await self.connect()
|
|
except:
|
|
pass
|
|
else:
|
|
self._bot.loop.create_task(self._handle_payload(msg.json()))
|
|
|
|
async def _handle_payload(self, data: dict) -> None:
|
|
op = data.get("op", None)
|
|
if not op:
|
|
return
|
|
|
|
if op == "ready":
|
|
self._session_id = data.get("sessionId")
|
|
|
|
if op == "stats":
|
|
self._stats = NodeStats(data)
|
|
return
|
|
|
|
if "guildId" in data:
|
|
if not (player := self._players.get(int(data["guildId"]))):
|
|
return
|
|
|
|
if op == "event":
|
|
await player._dispatch_event(data)
|
|
elif op == "playerUpdate":
|
|
await player._update_state(data)
|
|
|
|
async def send(self, method: RequestMethod, query: str, data: Union[dict, str] = {}) -> dict:
|
|
if not self._available:
|
|
raise NodeNotAvailable(f"The node '{self._identifier}' is unavailable.")
|
|
|
|
uri: str = f"{self._rest_uri}/{NODE_VERSION}/{query}"
|
|
async with self._session.request(
|
|
method=method.value,
|
|
url=uri,
|
|
headers={"Authorization": self._password},
|
|
json=data
|
|
) as resp:
|
|
if resp.status >= 300:
|
|
raise NodeException(f"Getting errors from Lavalink REST api")
|
|
|
|
if method == RequestMethod.DELETE:
|
|
return await resp.json(content_type=None)
|
|
|
|
return await resp.json()
|
|
|
|
async def connect(self) -> Node:
|
|
"""Initiates a connection with a Lavalink node and adds it to the node pool."""
|
|
|
|
try:
|
|
self._websocket = await self._session.ws_connect(
|
|
self._websocket_uri, headers=self._headers, heartbeat=self._heartbeat
|
|
)
|
|
|
|
self._task = self._bot.loop.create_task(self._listen())
|
|
self._available = True
|
|
self._info = NodeInfo(await self.send(RequestMethod.GET, query="info"))
|
|
|
|
self._logger.info(f"Node [{self._identifier}] is connected!")
|
|
|
|
except aiohttp.ClientConnectorError:
|
|
raise NodeConnectionFailure(
|
|
f"The connection to node '{self._identifier}' failed."
|
|
)
|
|
except aiohttp.WSServerHandshakeError:
|
|
raise NodeConnectionFailure(
|
|
f"The password for node '{self._identifier}' is invalid."
|
|
)
|
|
except aiohttp.InvalidURL:
|
|
raise NodeConnectionFailure(
|
|
f"The URI for node '{self._identifier}' is invalid."
|
|
)
|
|
|
|
if self.players:
|
|
await self.reconnect()
|
|
|
|
return self
|
|
|
|
async def disconnect(self, remove_from_pool: bool = False) -> None:
|
|
"""Disconnects a connected Lavalink node and removes it from the node pool.
|
|
This also destroys any players connected to the node.
|
|
"""
|
|
for player in self.players.copy().values():
|
|
await player.teardown()
|
|
|
|
await self._websocket.close()
|
|
if remove_from_pool:
|
|
del self._pool._nodes[self._identifier]
|
|
self._available = False
|
|
self._task.cancel()
|
|
|
|
self._logger.info(f"Node [{self._identifier}] is disconnected!")
|
|
|
|
async def reconnect(self) -> None:
|
|
await asyncio.sleep(10)
|
|
for player in self.players.copy().values():
|
|
try:
|
|
if player._voice_state:
|
|
await player._dispatch_voice_update(player._voice_state)
|
|
|
|
if player.current:
|
|
await player.play(track=player.current, start=min(player._last_position, player.current.length))
|
|
|
|
if player.is_paused:
|
|
await player.set_pause(True)
|
|
except:
|
|
await player.teardown()
|
|
await asyncio.sleep(2)
|
|
|
|
async def build_track(
|
|
self,
|
|
identifier: str,
|
|
requester: Member = None
|
|
) -> Track:
|
|
"""
|
|
Builds a track using a valid track identifier
|
|
|
|
You can also pass in a discord.py Context object to get a
|
|
Context object on the track it builds.
|
|
"""
|
|
|
|
async with self._session.get(
|
|
f"{self._rest_uri}/" + NODE_VERSION + "/decodetrack?",
|
|
headers={"Authorization": self._password},
|
|
params={"track": identifier}
|
|
) as resp:
|
|
if not resp.status == 200:
|
|
raise TrackLoadError(
|
|
f"Failed to build track. Check if the identifier is correct and try again."
|
|
)
|
|
|
|
data: dict = await resp.json()
|
|
return Track(track_id=identifier, info=data, requester=requester)
|
|
|
|
async def get_tracks(
|
|
self,
|
|
query: str,
|
|
*,
|
|
requester: Member,
|
|
search_type: SearchType = SearchType.YOUTUBE
|
|
) -> Union[List[Track], Playlist]:
|
|
"""Fetches tracks from the node's REST api to parse into Lavalink.
|
|
|
|
If you passed in Spotify API credentials, you can also pass in a
|
|
Spotify URL of a playlist, album or track and it will be parsed accordingly.
|
|
|
|
You can also pass in a discord.py Context object to get a
|
|
Context object on any track you search.
|
|
"""
|
|
|
|
if not URL_REGEX.match(query):
|
|
if search_type == SearchType.SPOTIFY:
|
|
return await self.spotifySearch(query=query, requester=requester)
|
|
|
|
else:
|
|
query = f"{search_type}:{query}"
|
|
|
|
if SPOTIFY_URL_REGEX.match(query):
|
|
try:
|
|
spotify_results = await self.spotify_client.search(query=query)
|
|
except Exception as _:
|
|
raise TrackLoadError("Not able to find the provided Spotify entity, is it private?")
|
|
|
|
if isinstance(spotify_results, spotify.Track):
|
|
return [
|
|
Track(
|
|
track_id=None,
|
|
info=spotify_results.to_dict(),
|
|
requester=requester,
|
|
search_type=search_type,
|
|
spotify_track=spotify_results,
|
|
)
|
|
]
|
|
|
|
tracks = [
|
|
Track(
|
|
track_id=None,
|
|
info=track.to_dict(),
|
|
requester=requester,
|
|
search_type=search_type,
|
|
spotify_track=track,
|
|
) for track in spotify_results.tracks if track.uri
|
|
]
|
|
|
|
return Playlist(
|
|
playlist_info={"name": spotify_results.name, "selectedTrack": 0},
|
|
tracks=tracks,
|
|
requester=requester,
|
|
spotify=True,
|
|
spotify_playlist=spotify_results
|
|
)
|
|
|
|
elif DISCORD_MP3_URL_REGEX.match(query):
|
|
async with self._session.get(
|
|
url=f"{self._rest_uri}/" + NODE_VERSION + f"/loadtracks?identifier={quote(query)}",
|
|
headers={"Authorization": self._password}
|
|
) as response:
|
|
data: dict = await response.json()
|
|
|
|
try:
|
|
track: dict = data["data"]
|
|
except:
|
|
raise TrackLoadError("Not able to find the provided track.")
|
|
|
|
return [
|
|
Track(
|
|
track_id=track["encoded"],
|
|
info=track["info"],
|
|
requester=requester
|
|
)
|
|
]
|
|
else:
|
|
async with self._session.get(
|
|
url=f"{self._rest_uri}/" + NODE_VERSION + f"/loadtracks?identifier={quote(query)}",
|
|
headers={"Authorization": self._password}
|
|
) as response:
|
|
data = await response.json()
|
|
|
|
load_type = data.get("loadType")
|
|
|
|
if not load_type:
|
|
raise TrackLoadError("There was an error while trying to load this track.")
|
|
|
|
elif load_type == "error":
|
|
exception = data["data"]
|
|
raise TrackLoadError(f"{exception['message']} [{exception['severity']}]")
|
|
|
|
elif load_type == "empty":
|
|
return None
|
|
|
|
elif load_type == "playlist":
|
|
data = data.get("data")
|
|
|
|
return Playlist(
|
|
playlist_info=data["info"],
|
|
tracks=data["tracks"],
|
|
requester=requester
|
|
)
|
|
|
|
elif load_type == "search":
|
|
return [
|
|
Track(
|
|
track_id=track["encoded"],
|
|
info=track["info"],
|
|
requester=requester
|
|
)
|
|
for track in data["data"]
|
|
]
|
|
|
|
elif load_type == "track":
|
|
track = data["data"]
|
|
return [
|
|
Track(
|
|
track_id=track["encoded"],
|
|
info=track["info"],
|
|
requester=requester
|
|
)
|
|
]
|
|
|
|
async def spotifySearch(self, query: str, *, requester: Member) -> Optional[List[Track]]:
|
|
try:
|
|
if not self.spotify_client:
|
|
raise InvalidSpotifyClientAuthorization(
|
|
"You did not provide proper Spotify client authorization credentials. "
|
|
"If you would like to use the Spotify searching feature, "
|
|
"please obtain Spotify API credentials here: https://developer.spotify.com/"
|
|
)
|
|
|
|
tracks = await self._spotify_client.track_search(query=query)
|
|
except Exception as _:
|
|
raise TrackLoadError("Not able to find the provided Spotify entity, is it private?")
|
|
|
|
return [
|
|
Track(
|
|
track_id=None,
|
|
requester=requester,
|
|
search_type=SearchType.YOUTUBE,
|
|
spotify_track=track,
|
|
info=track.to_dict()
|
|
)
|
|
for track in tracks ]
|
|
|
|
async def get_recommendations(self, track: Track, limit: int = 20) -> List[Optional[Track]]:
|
|
if track.spotify:
|
|
if not self.spotify_client:
|
|
return []
|
|
|
|
spotify_tracks = await self.spotify_client.similar_track(seed_tracks=track.identifier, limit=limit)
|
|
|
|
tracks = [
|
|
Track(
|
|
track_id=None,
|
|
search_type=SearchType.YOUTUBE,
|
|
spotify_track=track,
|
|
info=track.to_dict(),
|
|
requester=self.bot.user
|
|
)
|
|
for track in spotify_tracks
|
|
]
|
|
|
|
else:
|
|
if track.source != 'youtube':
|
|
return []
|
|
|
|
tracks = await self.get_tracks(
|
|
f"https://www.youtube.com/watch?v={track.identifier}&list=RD{track.identifier}",
|
|
requester=self.bot.user
|
|
)
|
|
|
|
if isinstance(tracks, Playlist):
|
|
tracks = tracks.tracks
|
|
|
|
return tracks[:limit] if limit else tracks
|
|
|
|
class NodePool:
|
|
"""The base class for the node pool.
|
|
This holds all the nodes that are to be used by the bot.
|
|
"""
|
|
|
|
_nodes: Dict[str, Node] = {}
|
|
|
|
def __repr__(self):
|
|
return f"<Voicelink.NodePool node_count={self.node_count}>"
|
|
|
|
@property
|
|
def nodes(self) -> Dict[str, Node]:
|
|
"""Property which returns a dict with the node identifier and the Node object."""
|
|
return self._nodes
|
|
|
|
@property
|
|
def node_count(self) -> Optional[Node]:
|
|
return len(self._nodes.values())
|
|
|
|
@classmethod
|
|
def get_best_node(cls, *, algorithm: NodeAlgorithm) -> Node:
|
|
"""Fetches the best node based on an NodeAlgorithm.
|
|
This option is preferred if you want to choose the best node
|
|
from a multi-node setup using either the node's latency
|
|
or the node's voice region.
|
|
Use NodeAlgorithm.by_ping if you want to get the best node
|
|
based on the node's latency.
|
|
Use NodeAlgorithm.by_region if you want to get the best node
|
|
based on the node's voice region. This method will only work
|
|
if you set a voice region when you create a node.
|
|
Use NodeAlgorithm.by_players if you want to get the best node
|
|
based on how players it has. This method will return a node with
|
|
the least amount of players
|
|
"""
|
|
available_nodes = [node for node in cls._nodes.values() if node._available]
|
|
|
|
if not available_nodes:
|
|
raise NoNodesAvailable("There are no nodes available.")
|
|
|
|
if algorithm == NodeAlgorithm.by_ping:
|
|
tested_nodes = {node: node.latency for node in available_nodes}
|
|
return min(tested_nodes, key=tested_nodes.get)
|
|
|
|
elif algorithm == NodeAlgorithm.by_players:
|
|
tested_nodes = {node: len(node.players.keys()) for node in available_nodes}
|
|
return min(tested_nodes, key=tested_nodes.get)
|
|
|
|
@classmethod
|
|
def get_node(cls, *, identifier: str = None) -> Node:
|
|
"""Fetches a node from the node pool using it's identifier.
|
|
If no identifier is provided, it will choose a node at random.
|
|
"""
|
|
|
|
available_nodes = { node
|
|
for _, node in cls._nodes.items() if node.is_connected
|
|
}
|
|
|
|
if identifier:
|
|
available_nodes = { node for node in available_nodes if node._identifier == identifier }
|
|
|
|
if not available_nodes:
|
|
raise NoNodesAvailable("There are no nodes available.")
|
|
|
|
nodes = {node: len(node.players.keys()) for node in available_nodes}
|
|
return min(nodes, key=nodes.get)
|
|
|
|
@classmethod
|
|
async def create_node(
|
|
cls,
|
|
*,
|
|
bot: Client,
|
|
host: str,
|
|
port: str,
|
|
password: str,
|
|
identifier: str,
|
|
secure: bool = False,
|
|
heartbeat: int = 30,
|
|
spotify_client_id: Optional[str] = None,
|
|
spotify_client_secret: Optional[str] = None,
|
|
session: Optional[aiohttp.ClientSession] = None,
|
|
resume_key: Optional[str] = None,
|
|
logger: Optional[logging.Logger] = None
|
|
) -> Node:
|
|
"""Creates a Node object to be then added into the node pool.
|
|
For Spotify searching capabilites, pass in valid Spotify API credentials.
|
|
"""
|
|
if identifier in cls._nodes.keys():
|
|
raise NodeCreationError(f"A node with identifier '{identifier}' already exists.")
|
|
|
|
if not logger:
|
|
logger = logging.getLogger("voicelink")
|
|
|
|
node = Node(
|
|
pool=cls, bot=bot, host=host, port=port, password=password,
|
|
identifier=identifier, secure=secure, heartbeat=heartbeat, spotify_client_id=spotify_client_id,
|
|
session=session, spotify_client_secret=spotify_client_secret,
|
|
resume_key=resume_key, logger=logger
|
|
)
|
|
|
|
await node.connect()
|
|
cls._nodes[node._identifier] = node
|
|
return node |