From 90c0c833885589df118c7ac88dc30082f5989c52 Mon Sep 17 00:00:00 2001 From: Choco <94597336+ChocoMeow@users.noreply.github.com> Date: Thu, 30 May 2024 21:42:05 +0800 Subject: [PATCH] Update methods.py --- ipc/methods.py | 184 ++++++++++++++++++++----------------------------- 1 file changed, 73 insertions(+), 111 deletions(-) diff --git a/ipc/methods.py b/ipc/methods.py index 70edcc5..1f59e0a 100644 --- a/ipc/methods.py +++ b/ipc/methods.py @@ -1,15 +1,14 @@ +import time import function as func -from typing import ( - List, - Dict, - Union -) +from typing import List, Dict, Union from discord import User, Member, VoiceChannel from discord.ext import commands from voicelink import Player, Track, Playlist, NodePool, decode, LoopType +RATELIMIT_COUNTER: Dict[int, Dict[str, float]] = {} + class TempCtx(): def __init__(self, author: Member, channel: VoiceChannel) -> None: self.author = author @@ -17,19 +16,20 @@ class TempCtx(): self.guild = channel.guild class SystemMethod: - def __init__(self, function): + def __init__(self, function: callable, *, credit: int = 1): self.function: callable = function self.params: List[str] = ["bot", "data"] + self.credit: int = credit class PlayerMethod(SystemMethod): - def __init__(self, function, *, auto_connect: bool = False): - super().__init__(function) + def __init__(self, function, *, credit: int = 1, auto_connect: bool = False): + super().__init__(function, credit=credit) self.params: List[str] = ["player", "member", "data"] self.auto_connect: bool = auto_connect class UserMethod(SystemMethod): - def __init__(self, function): - super().__init__(function) + def __init__(self, function, *, credit: int = 1): + super().__init__(function, credit=credit) self.params: List[str] = ["member", "data"] def missingPermission(user_id:int): @@ -46,7 +46,7 @@ def error_msg(msg: str, *, user_id: int = None, guild_id: int = None, level: str return payload -async def connect_channel(member: Member, bot: commands.Bot): +async def connect_channel(member: Member, bot: commands.Bot) -> Player: if not member.voice: return @@ -84,7 +84,7 @@ async def initUser(bot: commands.Bot, data: Dict) -> Dict: "data": data } -async def initPlayer(player: Player, member: Member, data: Dict): +async def initPlayer(player: Player, member: Member, data: Dict) -> Dict: player._ipc_connection = True return { "op": "initPlayer", @@ -98,7 +98,7 @@ async def initPlayer(player: Player, member: Member, data: Dict): "tracks": [ track.track_id for track in player.queue._queue ], "repeat_mode": player.queue.repeat.lower(), "channel_name": player.channel.name, - "current_queue_position": player.queue._position if player._current else player.queue._position + 1, + "current_queue_position": player.queue._position if player._current else player.queue._position, "current_position": 0 or player.position if player.is_playing else 0, "is_playing": player.is_playing, "is_paused": player.is_paused, @@ -106,7 +106,10 @@ async def initPlayer(player: Player, member: Member, data: Dict): "autoplay": player.settings.get("autoplay", False) } -async def getRecommendation(bot: commands.Bot, data: Dict): +async def closeConnection(player: Player, member: Member, data: Dict) -> None: + player._ipc_connection = False + +async def getRecommendation(bot: commands.Bot, data: Dict) -> None: node = NodePool.get_node() if not node: return @@ -122,7 +125,7 @@ async def getRecommendation(bot: commands.Bot, data: Dict): "tracks": [track.track_id for track in tracks] if tracks else [] } -async def skipTo(player: Player, member: Member, data: Dict): +async def skipTo(player: Player, member: Member, data: Dict) -> None: if not player.is_privileged(member): if player.current and member == player.current.requester: pass @@ -144,7 +147,7 @@ async def skipTo(player: Player, member: Member, data: Dict): await player.set_repeat(LoopType.off.name) await player.stop() -async def backTo(player: Player, member: Member, data: Dict): +async def backTo(player: Player, member: Member, data: Dict) -> None: if not player.is_privileged(member): if player.current and member == player.current.requester: pass @@ -177,20 +180,29 @@ async def moveTrack(player: Player, member: Member, data: Dict) -> None: await player.move_track(index, new_index, member) -async def addTracks(player: Player, member: Member, data: Dict): - raw_tracks = data.get("tracks", []) +async def addTracks(player: Player, member: Member, data: Dict) -> None: + _type = data.get("type", "addToQueue") tracks = [Track( track_id=track_id, info=decode(track_id), requester=member - ) for track_id in raw_tracks] + ) for track_id in data.get("tracks", [])] - await player.add_track(tracks) + if _type == "addToQueue": + await player.add_track(tracks) + + elif _type == "forcePlay": + await player.add_track(tracks, at_font=True) + if player.is_playing: + return await player.stop() + + elif _type == "addNext": + await player.add_track(tracks, at_font=True) if not player.is_playing: await player.do_next() -async def getTracks(player: Player, member: Member, data: Dict): +async def getTracks(player: Player, member: Member, data: Dict) -> Dict: query = data.get("query", None) if query: @@ -205,7 +217,7 @@ async def getTracks(player: Player, member: Member, data: Dict): payload["tracks"] = [ track.track_id for track in tracks ] return payload -async def shuffleTrack(player: Player, member: Member, data: Dict): +async def shuffleTrack(player: Player, member: Member, data: Dict) -> None: if not player.is_privileged(member): if member in player.shuffle_votes: @@ -219,7 +231,7 @@ async def shuffleTrack(player: Player, member: Member, data: Dict): await player.shuffle(data.get("type", "queue"), member) -async def repeatTrack(player: Player, member: Member, data: Dict): +async def repeatTrack(player: Player, member: Member, data: Dict) -> None: if not player.is_privileged(member): return missingPermission(member.id) @@ -232,7 +244,7 @@ async def removeTrack(player: Player, member: Member, data: Dict) -> None: index, index2 = data.get("index"), data.get("index2") await player.remove_track(index, index2, requester=member) -async def updatePause(player: Player, member: Member, data: Dict): +async def updatePause(player: Player, member: Member, data: Dict) -> None: pause = data.get("pause", True) if not player.is_privileged(member): if pause: @@ -253,22 +265,18 @@ async def updatePause(player: Player, member: Member, data: Dict): pass else: return error_msg(player.get_msg('resumeVote').format(member, len(player.resume_votes), required), guild_id=player.guild.id) - - if pause: - player.pause_votes.clear() - else: - player.resume_votes.clear() + player.pause_votes.clear() if pause else player.resume_votes.clear() await player.set_pause(pause, member) -async def updatePosition(player: Player, member: Member, data: Dict): +async def updatePosition(player: Player, member: Member, data: Dict) -> None: if not player.is_privileged(member): return missingPermission(member.id) position = data.get("position"); await player.seek(position, member); -async def toggleAutoplay(player: Player, member: Member, data: Dict): +async def toggleAutoplay(player: Player, member: Member, data: Dict) -> Dict: if not player.is_privileged(member): return error_msg(player.get_msg('missingPerms_autoplay')) @@ -285,91 +293,40 @@ async def toggleAutoplay(player: Player, member: Member, data: Dict): "requester_id": member.id } -async def closeConnection(player: Player, member: Member, data: Dict): - player._ipc_connection = False +async def _loadPlaylist(playlist_id: str, playlist: Dict, user_id: int) -> Dict: + payload = {"op": "loadPlaylist", "playlist_id": playlist_id, "user_id": user_id} -async def getPlaylists(member: Member, data: Dict): - playlists: Dict = await func.get_user(member.id, "playlist") - if not playlists: - return + if playlist.get("type") == "link": + tracks: List[Track]= await NodePool.get_node().get_tracks(playlist.get("uri"), requester=None) + if tracks: + payload["tracks"] = [track.track_id for track in (tracks.tracks if isinstance(tracks, Playlist) else tracks)] + else: + payload["tracks"] = playlist.get("tracks", []) - for pId, pList in playlists.copy().items(): - if "type" in pList: - if pList["type"] == "link": - tracks: Playlist = await NodePool.get_node().get_tracks(pList["uri"], requester=member) - if tracks: - playlists[pId]["tracks"] = [ track.track_id for track in tracks.tracks ] + return payload - elif pList["type"] == "share": - playlist = await func.get_user(pList["user"], "playlist") - playlist = playlist.get(pList["referId"]) - if playlist: - if member.id not in playlist["perms"]["read"]: - await func.update_user(member.id, {"$unset": {f"playlist.{pId}": 1}}) - del playlists[pId] - continue - - if playlist['type'] == 'link': - tracks: Playlist = await NodePool.get_node().get_tracks(playlist["uri"], requester=member) - playlists[pId]["tracks"] = [ track.track_id for track in tracks.tracks ] - else: - playlists[pId]["tracks"] = playlist["tracks"] - - return { - "op": "getPlaylists", - "playlists": playlists, - "user_id": member.id - } - -async def removePlaylist(member: Member, data: Dict): - pId = data.get("pId") - isShare = data.get("isShare", False) - - if pId == 200: - return - - if isShare: - refer_user = data.get("refer_user") - await func.update_user(refer_user, {"$pull": {f"playlist.{pId}.perms.read": member.id}}) - - await func.update_user(member.id, {"$unset": {f'playlist.{pId}': 1}}) - -async def addPlaylistTrack(member: Member, data: Dict): - track_id = data.get("track_id") - pId = data.get("pId") - if not track_id or not pId: - return +async def getPlaylist(bot: commands.Bot, data: Dict) -> None: + user_id = data.get("user_id") + playlist_id = str(data.get("playlist_id")) + playlists = await func.get_user(user_id, "playlist") - playlist: Dict = await func.get_user(member.id, 'playlist') - playlist = playlist.get(pId) + playlist = playlists.get(playlist_id) if not playlist: return - if playlist["type"] != "playlist": - return error_msg(func.get_lang(member.guild.id, 'playlistNotAllow'), user_id=member.id) - - rank, max_p, max_t = func.check_roles() - if len(playlist["tracks"]) >= max_t: - return error_msg(func.get_lang(member.guild.id, "playlistlimited").format(max_t), user_id=member.id) - - if track_id in playlist['tracks']: - return error_msg(func.get_lang(member.guild.id, "playlistrepeated"), user_id=member.id) - - await func.update_user(member.id, {"$push": {f'playlist.{pId}.tracks': track_id}}) - -async def removePlaylistTrack(member: Member, data: Dict): - track_id = data.get("track_id") - pId = data.get("pId") - if not track_id or not pId: - return - - await func.update_user(member.id, {"$pull": {f'playlist.{pId}.tracks': track_id }}) + if playlist["type"] == "share": + target_user = await func.get_user(playlist["user"], "playlist") + target_playlist = target_user.get(playlist["referId"]) + if target_playlist and user_id in target_playlist["perms"]["read"]: + return await _loadPlaylist(playlist_id, target_playlist, user_id) + + return await _loadPlaylist(playlist_id, playlist, user_id) methods: Dict[str, Union[SystemMethod, PlayerMethod, UserMethod]] = { - "initBot": SystemMethod(initBot), - "initUser": SystemMethod(initUser), - "getRecommendation": SystemMethod(getRecommendation), - "closeConnection": PlayerMethod(closeConnection), + "initBot": SystemMethod(initBot, credit=0), + "initUser": SystemMethod(initUser, credit=0), + "getRecommendation": SystemMethod(getRecommendation, credit=4), + "closeConnection": PlayerMethod(closeConnection, credit=0), "initPlayer": PlayerMethod(initPlayer), "skipTo": PlayerMethod(skipTo), "backTo": PlayerMethod(backTo), @@ -382,10 +339,7 @@ methods: Dict[str, Union[SystemMethod, PlayerMethod, UserMethod]] = { "updatePause": PlayerMethod(updatePause), "updatePosition": PlayerMethod(updatePosition), "toggleAutoplay": PlayerMethod(toggleAutoplay), - "getPlaylists": UserMethod(getPlaylists), - "removePlaylist": UserMethod(removePlaylist), - "addPlaylistTrack": UserMethod(addPlaylistTrack), - "removePlaylistTrack": UserMethod(removePlaylistTrack) + "getPlaylist": SystemMethod(getPlaylist) } async def process_methods(ipc_client, bot: commands.Bot, data: Dict) -> None: @@ -394,6 +348,14 @@ async def process_methods(ipc_client, bot: commands.Bot, data: Dict) -> None: if not method or not (user_id := data.get("user_id")): return + if user_id not in RATELIMIT_COUNTER or (time.time() - RATELIMIT_COUNTER[user_id]["time"]) >= 300: + RATELIMIT_COUNTER[user_id] = {"time": time.time(), "count": 0} + + else: + if RATELIMIT_COUNTER[user_id]["count"] >= 200: + return await ipc_client.send({"op": "rateLimited", "user_id": user_id}) + RATELIMIT_COUNTER[user_id]["count"] += method.credit + try: env: Dict = {"bot": bot, "data": data} args: List = []