From 655f41d1c5bcdfe58acb5f4a726980d52f24d085 Mon Sep 17 00:00:00 2001 From: jack1142 <6032823+jack1142@users.noreply.github.com> Date: Mon, 18 Jan 2021 13:01:47 +0100 Subject: [PATCH 1/8] Use partial messages in Streams cog to avoid leakage --- redbot/cogs/streams/streams.py | 51 ++++++++++++++++-------------- redbot/cogs/streams/streamtypes.py | 34 +++++++++++++------- 2 files changed, 50 insertions(+), 35 deletions(-) diff --git a/redbot/cogs/streams/streams.py b/redbot/cogs/streams/streams.py index f490bc5ce72..b54a098a7fe 100644 --- a/redbot/cogs/streams/streams.py +++ b/redbot/cogs/streams/streams.py @@ -714,14 +714,23 @@ async def _stream_alerts(self): await asyncio.sleep(await self.config.refresh_timer()) async def _send_stream_alert( - self, stream, channel: discord.TextChannel, embed: discord.Embed, content: str = None + self, + stream, + channel: discord.TextChannel, + embed: discord.Embed, + content: str = None, + *, + is_schedule: bool = False, ): m = await channel.send( content, embed=embed, allowed_mentions=discord.AllowedMentions(roles=True, everyone=True), ) - stream._messages_cache.append(m) + message_data = {"channel": m.channel.id, "message": m.id} + if is_schedule: + message_data["is_schedule"] = True + stream.messages.append(message_data) async def check_streams(self): for stream in self.streams: @@ -739,19 +748,25 @@ async def check_streams(self): else: embed = await stream.is_online() except OfflineStream: - if not stream._messages_cache: + if not stream.messages: continue - for message in stream._messages_cache: - if await self.bot.cog_disabled_in_guild(self, message.guild): + + for msg_data in stream.iter_messages(): + partial_msg = msg_data["partial_message"] + if partial_msg is None: + continue + if await self.bot.cog_disabled_in_guild(self, partial_msg.guild): continue - autodelete = await self.config.guild(message.guild).autodelete() - if autodelete: - with contextlib.suppress(discord.NotFound): - await message.delete() - stream._messages_cache.clear() + if not await self._config.guild(partial_msg.guild).autodelete(): + continue + + with contextlib.suppress(discord.NotFound): + await partial_msg.delete() + + stream.messages.clear() await self.save_streams() else: - if stream._messages_cache: + if stream.messages: continue for channel_id in stream.channels: channel = self.bot.get_channel(channel_id) @@ -767,7 +782,7 @@ async def check_streams(self): continue if is_schedule: # skip messages and mentions - await self._send_stream_alert(stream, channel, embed) + await self._send_stream_alert(stream, channel, embed, is_schedule=True) await self.save_streams() continue await set_contextual_locales_from_guild(self.bot, channel.guild) @@ -864,17 +879,6 @@ async def load_streams(self): _class = getattr(_streamtypes, raw_stream["type"], None) if not _class: continue - raw_msg_cache = raw_stream["messages"] - raw_stream["_messages_cache"] = [] - for raw_msg in raw_msg_cache: - chn = self.bot.get_channel(raw_msg["channel"]) - if chn is not None: - try: - msg = await chn.fetch_message(raw_msg["message"]) - except discord.HTTPException: - pass - else: - raw_stream["_messages_cache"].append(msg) token = await self.bot.get_shared_api_tokens(_class.token_name) if token: if _class.__name__ == "TwitchStream": @@ -884,6 +888,7 @@ async def load_streams(self): if _class.__name__ == "YoutubeStream": raw_stream["config"] = self.config raw_stream["token"] = token + raw_stream["bot"] = self.bot streams.append(_class(**raw_stream)) return streams diff --git a/redbot/cogs/streams/streamtypes.py b/redbot/cogs/streams/streamtypes.py index c1e315e5928..a24c3234e21 100644 --- a/redbot/cogs/streams/streamtypes.py +++ b/redbot/cogs/streams/streamtypes.py @@ -56,10 +56,11 @@ class Stream: token_name: ClassVar[Optional[str]] = None def __init__(self, **kwargs): + self.bot = kwargs.pop("bot") self.name = kwargs.pop("name", None) self.channels = kwargs.pop("channels", []) # self.already_online = kwargs.pop("already_online", False) - self._messages_cache = kwargs.pop("_messages_cache", []) + self.messages = kwargs.pop("messages", []) self.type = self.__class__.__name__ async def is_online(self): @@ -68,14 +69,19 @@ async def is_online(self): def make_embed(self): raise NotImplementedError() + def iter_messages(self): + for msg_data in self.messages: + data = msg_data.copy() + channel = self.bot.get_channel(msg_data["channel"]) + if channel is not None: + data["partial_message"] = channel.get_partial_message(data["message"]) + yield data + def export(self): data = {} for k, v in self.__dict__.items(): if not k.startswith("_"): data[k] = v - data["messages"] = [] - for m in self._messages_cache: - data["messages"].append({"channel": m.channel.id, "message": m.id}) return data def __repr__(self): @@ -190,17 +196,21 @@ async def make_embed(self, data): embed.timestamp = start_time is_schedule = True else: - # repost message + # delete the message(s) about the stream schedule to_remove = [] - for message in self._messages_cache: - if message.embeds[0].description is discord.Embed.Empty: + for msg_data in self.iter_messages(): + if not msg_data.get("is_schedule", False): continue - with contextlib.suppress(Exception): - autodelete = await self._config.guild(message.guild).autodelete() + partial_msg = msg_data["partial_message"] + if partial_msg is not None: + autodelete = await self._config.guild(partial_msg.guild).autodelete() if autodelete: - await message.delete() - to_remove.append(message.id) - self._messages_cache = [x for x in self._messages_cache if x.id not in to_remove] + with contextlib.suppress(discord.NotFound): + await partial_msg.delete() + to_remove.append(msg_data["message"]) + self.messages = [ + data for data in self.messages if data["message"] not in to_remove + ] embed.set_author(name=channel_title) embed.set_image(url=rnd(thumbnail)) embed.colour = 0x9255A5 From 0dba2b6f15ecaaa7dc05c5fb5030cfeb896c3dd0 Mon Sep 17 00:00:00 2001 From: jack1142 <6032823+jack1142@users.noreply.github.com> Date: Mon, 5 Apr 2021 19:32:33 +0200 Subject: [PATCH 2/8] Stop trying to save bot object to Config... --- redbot/cogs/streams/streams.py | 2 +- redbot/cogs/streams/streamtypes.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/redbot/cogs/streams/streams.py b/redbot/cogs/streams/streams.py index b54a098a7fe..91f47f8c492 100644 --- a/redbot/cogs/streams/streams.py +++ b/redbot/cogs/streams/streams.py @@ -888,7 +888,7 @@ async def load_streams(self): if _class.__name__ == "YoutubeStream": raw_stream["config"] = self.config raw_stream["token"] = token - raw_stream["bot"] = self.bot + raw_stream["_bot"] = self.bot streams.append(_class(**raw_stream)) return streams diff --git a/redbot/cogs/streams/streamtypes.py b/redbot/cogs/streams/streamtypes.py index a24c3234e21..a9ca3eb4f56 100644 --- a/redbot/cogs/streams/streamtypes.py +++ b/redbot/cogs/streams/streamtypes.py @@ -56,7 +56,7 @@ class Stream: token_name: ClassVar[Optional[str]] = None def __init__(self, **kwargs): - self.bot = kwargs.pop("bot") + self._bot = kwargs.pop("_bot") self.name = kwargs.pop("name", None) self.channels = kwargs.pop("channels", []) # self.already_online = kwargs.pop("already_online", False) @@ -72,7 +72,7 @@ def make_embed(self): def iter_messages(self): for msg_data in self.messages: data = msg_data.copy() - channel = self.bot.get_channel(msg_data["channel"]) + channel = self._bot.get_channel(msg_data["channel"]) if channel is not None: data["partial_message"] = channel.get_partial_message(data["message"]) yield data From 4d920b9e6d869f0ce9d2d2404f67e6faf908bf63 Mon Sep 17 00:00:00 2001 From: jack1142 <6032823+jack1142@users.noreply.github.com> Date: Mon, 5 Apr 2021 20:00:33 +0200 Subject: [PATCH 3/8] Put guild id as part of message data --- redbot/cogs/streams/streams.py | 2 +- redbot/cogs/streams/streamtypes.py | 7 ++++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/redbot/cogs/streams/streams.py b/redbot/cogs/streams/streams.py index 91f47f8c492..15fc15106b9 100644 --- a/redbot/cogs/streams/streams.py +++ b/redbot/cogs/streams/streams.py @@ -727,7 +727,7 @@ async def _send_stream_alert( embed=embed, allowed_mentions=discord.AllowedMentions(roles=True, everyone=True), ) - message_data = {"channel": m.channel.id, "message": m.id} + message_data = {"guild": m.guild.id, "channel": m.channel.id, "message": m.id} if is_schedule: message_data["is_schedule"] = True stream.messages.append(message_data) diff --git a/redbot/cogs/streams/streamtypes.py b/redbot/cogs/streams/streamtypes.py index a9ca3eb4f56..3be9686cdc2 100644 --- a/redbot/cogs/streams/streamtypes.py +++ b/redbot/cogs/streams/streamtypes.py @@ -72,7 +72,12 @@ def make_embed(self): def iter_messages(self): for msg_data in self.messages: data = msg_data.copy() - channel = self._bot.get_channel(msg_data["channel"]) + # "guild" key might not exist for old config data (available since GH-4742) + if guild_id := msg_data.get("guild"): + guild = self._bot.get_guild(guild_id) + channel = guild and guild.get_channel(msg_data["channel"]) + else: + channel = self._bot.get_channel(msg_data["channel"]) if channel is not None: data["partial_message"] = channel.get_partial_message(data["message"]) yield data From 7c30fac20d48e30d2255e744ba292bd21550203a Mon Sep 17 00:00:00 2001 From: jack1142 <6032823+jack1142@users.noreply.github.com> Date: Mon, 5 Apr 2021 20:38:27 +0200 Subject: [PATCH 4/8] Fix AttributeError --- redbot/cogs/streams/streams.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/redbot/cogs/streams/streams.py b/redbot/cogs/streams/streams.py index 15fc15106b9..fb68487f686 100644 --- a/redbot/cogs/streams/streams.py +++ b/redbot/cogs/streams/streams.py @@ -757,7 +757,7 @@ async def check_streams(self): continue if await self.bot.cog_disabled_in_guild(self, partial_msg.guild): continue - if not await self._config.guild(partial_msg.guild).autodelete(): + if not await self.config.guild(partial_msg.guild).autodelete(): continue with contextlib.suppress(discord.NotFound): From 530fee90296371f8914d8ebbdbfe7eedfdf58c05 Mon Sep 17 00:00:00 2001 From: jack1142 <6032823+jack1142@users.noreply.github.com> Date: Mon, 5 Apr 2021 20:40:04 +0200 Subject: [PATCH 5/8] Pass bot object to stream classes in commands --- redbot/cogs/streams/streams.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/redbot/cogs/streams/streams.py b/redbot/cogs/streams/streams.py index fb68487f686..bbcdf184ebe 100644 --- a/redbot/cogs/streams/streams.py +++ b/redbot/cogs/streams/streams.py @@ -209,6 +209,7 @@ async def twitchstream(self, ctx: commands.Context, channel_name: str): await self.maybe_renew_twitch_bearer_token() token = (await self.bot.get_shared_api_tokens("twitch")).get("client_id") stream = TwitchStream( + bot=self.bot, name=channel_name, token=token, bearer=self.ttv_bearer_cache.get("access_token", None), @@ -224,21 +225,25 @@ async def youtubestream(self, ctx: commands.Context, channel_id_or_name: str): apikey = await self.bot.get_shared_api_tokens("youtube") is_name = self.check_name_or_id(channel_id_or_name) if is_name: - stream = YoutubeStream(name=channel_id_or_name, token=apikey, config=self.config) + stream = YoutubeStream( + bot=self.bot, name=channel_id_or_name, token=apikey, config=self.config + ) else: - stream = YoutubeStream(id=channel_id_or_name, token=apikey, config=self.config) + stream = YoutubeStream( + bot=self.bot, id=channel_id_or_name, token=apikey, config=self.config + ) await self.check_online(ctx, stream) @commands.command() async def smashcast(self, ctx: commands.Context, channel_name: str): """Check if a smashcast channel is live.""" - stream = HitboxStream(name=channel_name) + stream = HitboxStream(bot=self.bot, name=channel_name) await self.check_online(ctx, stream) @commands.command() async def picarto(self, ctx: commands.Context, channel_name: str): """Check if a Picarto channel is live.""" - stream = PicartoStream(name=channel_name) + stream = PicartoStream(bot=self.bot, name=channel_name) await self.check_online(ctx, stream) async def check_online( From ba3e516adcf84310daf44785e3fdc4c71a48006b Mon Sep 17 00:00:00 2001 From: jack1142 <6032823+jack1142@users.noreply.github.com> Date: Mon, 5 Apr 2021 20:44:33 +0200 Subject: [PATCH 6/8] ugh --- redbot/cogs/streams/streams.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/redbot/cogs/streams/streams.py b/redbot/cogs/streams/streams.py index bbcdf184ebe..11274a8f8a0 100644 --- a/redbot/cogs/streams/streams.py +++ b/redbot/cogs/streams/streams.py @@ -209,7 +209,7 @@ async def twitchstream(self, ctx: commands.Context, channel_name: str): await self.maybe_renew_twitch_bearer_token() token = (await self.bot.get_shared_api_tokens("twitch")).get("client_id") stream = TwitchStream( - bot=self.bot, + _bot=self.bot, name=channel_name, token=token, bearer=self.ttv_bearer_cache.get("access_token", None), @@ -226,24 +226,24 @@ async def youtubestream(self, ctx: commands.Context, channel_id_or_name: str): is_name = self.check_name_or_id(channel_id_or_name) if is_name: stream = YoutubeStream( - bot=self.bot, name=channel_id_or_name, token=apikey, config=self.config + _bot=self.bot, name=channel_id_or_name, token=apikey, config=self.config ) else: stream = YoutubeStream( - bot=self.bot, id=channel_id_or_name, token=apikey, config=self.config + _bot=self.bot, id=channel_id_or_name, token=apikey, config=self.config ) await self.check_online(ctx, stream) @commands.command() async def smashcast(self, ctx: commands.Context, channel_name: str): """Check if a smashcast channel is live.""" - stream = HitboxStream(bot=self.bot, name=channel_name) + stream = HitboxStream(_bot=self.bot, name=channel_name) await self.check_online(ctx, stream) @commands.command() async def picarto(self, ctx: commands.Context, channel_name: str): """Check if a Picarto channel is live.""" - stream = PicartoStream(bot=self.bot, name=channel_name) + stream = PicartoStream(_bot=self.bot, name=channel_name) await self.check_online(ctx, stream) async def check_online( From 4b969b12cbf9ae4edbff7f0f2c8d53ea50a50db7 Mon Sep 17 00:00:00 2001 From: jack1142 <6032823+jack1142@users.noreply.github.com> Date: Mon, 5 Apr 2021 21:03:39 +0200 Subject: [PATCH 7/8] Another place we use this class in --- redbot/cogs/streams/streams.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/redbot/cogs/streams/streams.py b/redbot/cogs/streams/streams.py index 11274a8f8a0..cb2c62446eb 100644 --- a/redbot/cogs/streams/streams.py +++ b/redbot/cogs/streams/streams.py @@ -401,7 +401,7 @@ async def stream_alert(self, ctx: commands.Context, _class, channel_name): is_yt = _class.__name__ == "YoutubeStream" is_twitch = _class.__name__ == "TwitchStream" if is_yt and not self.check_name_or_id(channel_name): - stream = _class(id=channel_name, token=token, config=self.config) + stream = _class(_bot=self.bot, id=channel_name, token=token, config=self.config) elif is_twitch: await self.maybe_renew_twitch_bearer_token() stream = _class( From 1d0b7bfcf6539b95c2e8da118fa5dcc420403d59 Mon Sep 17 00:00:00 2001 From: jack1142 <6032823+jack1142@users.noreply.github.com> Date: Mon, 5 Apr 2021 21:24:13 +0200 Subject: [PATCH 8/8] more... --- redbot/cogs/streams/streams.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/redbot/cogs/streams/streams.py b/redbot/cogs/streams/streams.py index cb2c62446eb..594f0969af2 100644 --- a/redbot/cogs/streams/streams.py +++ b/redbot/cogs/streams/streams.py @@ -405,15 +405,18 @@ async def stream_alert(self, ctx: commands.Context, _class, channel_name): elif is_twitch: await self.maybe_renew_twitch_bearer_token() stream = _class( + _bot=self.bot, name=channel_name, token=token.get("client_id"), bearer=self.ttv_bearer_cache.get("access_token", None), ) else: if is_yt: - stream = _class(name=channel_name, token=token, config=self.config) + stream = _class( + _bot=self.bot, name=channel_name, token=token, config=self.config + ) else: - stream = _class(name=channel_name, token=token) + stream = _class(_bot=self.bot, name=channel_name, token=token) try: exists = await self.check_exists(stream) except InvalidTwitchCredentials: