diff options
| -rw-r--r-- | bot/cogs/bigbrother.py | 96 | ||||
| -rw-r--r-- | bot/constants.py | 8 | ||||
| -rw-r--r-- | bot/utils/messages.py | 40 | ||||
| -rw-r--r-- | config-default.yml | 6 |
4 files changed, 136 insertions, 14 deletions
diff --git a/bot/cogs/bigbrother.py b/bot/cogs/bigbrother.py index 3f30eb0e9..68ae4546b 100644 --- a/bot/cogs/bigbrother.py +++ b/bot/cogs/bigbrother.py @@ -1,16 +1,21 @@ +import asyncio import logging +import re +from collections import defaultdict, deque from typing import List, Union from discord import Color, Embed, Guild, Member, Message, TextChannel, User from discord.ext.commands import Bot, Context, group -from bot.constants import Channels, Emojis, Guild as GuildConfig, Keys, Roles, URLs +from bot.constants import BigBrother as BigBrotherConfig, Channels, Emojis, Guild as GuildConfig, Keys, Roles, URLs from bot.decorators import with_role from bot.pagination import LinePaginator - +from bot.utils import messages log = logging.getLogger(__name__) +URL_RE = re.compile(r"(https?://[^\s]+)") + class BigBrother: """User monitoring to assist with moderation.""" @@ -19,7 +24,11 @@ class BigBrother: def __init__(self, bot: Bot): self.bot = bot - self.watched_users = {} + self.watched_users = {} # { user_id: log_channel_id } + self.channel_queues = defaultdict(lambda: defaultdict(deque)) # { user_id: { channel_id: queue(messages) } + self.consuming = False + + self.bot.loop.create_task(self.get_watched_users()) def update_cache(self, api_response: List[dict]): """ @@ -43,7 +52,10 @@ class BigBrother: "but the given channel could not be found. Ignoring." ) - async def on_ready(self): + async def get_watched_users(self): + """Retrieves watched users from the API.""" + + await self.bot.wait_until_ready() async with self.bot.http_session.get(URLs.site_bigbrother_api, headers=self.HEADERS) as response: data = await response.json() self.update_cache(data) @@ -55,9 +67,10 @@ class BigBrother: async with self.bot.http_session.delete(url, headers=self.HEADERS) as response: del self.watched_users[user.id] + del self.channel_queues[user.id] if response.status == 204: await channel.send( - f"{Emojis.lemoneye2}:hammer: {user} got banned, so " + f"{Emojis.bb_message}:hammer: {user} got banned, so " f"`BigBrother` will no longer relay their messages to {channel}" ) @@ -65,19 +78,77 @@ class BigBrother: data = await response.json() reason = data.get('error_message', "no message provided") await channel.send( - f"{Emojis.lemoneye2}:x: {user} got banned, but trying to remove them from" + f"{Emojis.bb_message}:x: {user} got banned, but trying to remove them from" f"BigBrother's user dictionary on the API returned an error: {reason}" ) async def on_message(self, msg: Message): + """Queues up messages sent by watched users.""" + if msg.author.id in self.watched_users: - channel = self.watched_users[msg.author.id] - relay_content = (f"{Emojis.lemoneye2} {msg.author} sent the following " - f"in {msg.channel.mention}: {msg.clean_content}") - if msg.attachments: - relay_content += f" (with {len(msg.attachments)} attachment(s))" + if not self.consuming: + self.bot.loop.create_task(self.consume_messages()) + + log.trace(f"Received message: {msg.content} ({len(msg.attachments)} attachments)") + self.channel_queues[msg.author.id][msg.channel.id].append(msg) + + async def consume_messages(self): + """Consumes the message queues to log watched users' messages.""" + + if not self.consuming: + self.consuming = True + log.trace("Sleeping before consuming...") + await asyncio.sleep(BigBrotherConfig.log_delay) + + log.trace("Begin consuming messages.") + channel_queues = self.channel_queues.copy() + self.channel_queues.clear() + for user_id, queues in channel_queues.items(): + for _, queue in queues.items(): + channel = self.watched_users[user_id] + + if queue: + # Send a header embed before sending all messages in the queue. + msg = queue[0] + embed = Embed(description=f"{msg.author.mention} in [#{msg.channel.name}]({msg.jump_url})") + embed.set_author(name=msg.author.nick or msg.author.name, icon_url=msg.author.avatar_url) + await channel.send(embed=embed) + + while queue: + msg = queue.popleft() + log.trace(f"Consuming message: {msg.clean_content} ({len(msg.attachments)} attachments)") + await self.log_message(msg, channel) + + if self.channel_queues: + log.trace("Queue not empty; continue consumption.") + self.bot.loop.create_task(self.consume_messages()) + else: + log.trace("Done consuming messages.") + self.consuming = False + + @staticmethod + async def log_message(message: Message, destination: TextChannel): + """ + Logs a watched user's message in the given channel. + + Attachments are also sent. All non-image or non-video URLs are put in inline code blocks to prevent preview + embeds from being automatically generated. + + :param message: the message to log + :param destination: the channel in which to log the message + """ + + content = message.clean_content + if content: + # Put all non-media URLs in inline code blocks. + media_urls = {embed.url for embed in message.embeds if embed.type in ("image", "video")} + for url in URL_RE.findall(content): + if url not in media_urls: + content = content.replace(url, f"`{url}`") + + await destination.send(content) - await channel.send(relay_content) + await messages.send_attachments(message, destination) @group(name='bigbrother', aliases=('bb',), invoke_without_command=True) @with_role(Roles.owner, Roles.admin, Roles.moderator) @@ -175,6 +246,7 @@ class BigBrother: if user.id in self.watched_users: del self.watched_users[user.id] + del self.channel_queues[user.id] else: log.warning(f"user {user.id} was unwatched but was not found in the cache") diff --git a/bot/constants.py b/bot/constants.py index 2a9796cc5..0e8c52c68 100644 --- a/bot/constants.py +++ b/bot/constants.py @@ -241,7 +241,7 @@ class Emojis(metaclass=YAMLGetter): green_chevron: str red_chevron: str white_chevron: str - lemoneye2: str + bb_message: str status_online: str status_offline: str @@ -446,6 +446,12 @@ class AntiSpam(metaclass=YAMLGetter): rules: Dict[str, Dict[str, int]] +class BigBrother(metaclass=YAMLGetter): + section = 'big_brother' + + log_delay: int + + # Debug mode DEBUG_MODE = True if 'local' in os.environ.get("SITE_URL", "local") else False diff --git a/bot/utils/messages.py b/bot/utils/messages.py index c625beb5c..63e41983b 100644 --- a/bot/utils/messages.py +++ b/bot/utils/messages.py @@ -1,9 +1,13 @@ import asyncio import contextlib +from io import BytesIO from typing import Sequence -from discord import Message +from discord import Embed, File, Message, TextChannel from discord.abc import Snowflake +from discord.errors import HTTPException + +MAX_SIZE = 1024 * 1024 * 8 # 8 Mebibytes async def wait_for_deletion( @@ -70,3 +74,37 @@ async def wait_for_deletion( timeout=timeout ) await message.delete() + + +async def send_attachments(message: Message, destination: TextChannel): + """ + Re-uploads each attachment in a message to the given channel. + + Each attachment is sent as a separate message to more easily comply with the 8 MiB request size limit. + If attachments are too large, they are instead grouped into a single embed which links to them. + + :param message: the message whose attachments to re-upload + :param destination: the channel in which to re-upload the attachments + """ + + large = [] + for attachment in message.attachments: + try: + # This should avoid most files that are too large, but some may get through hence the try-catch. + # Allow 512 bytes of leeway for the rest of the request. + if attachment.size <= MAX_SIZE - 512: + with BytesIO() as file: + await attachment.save(file) + await destination.send(file=File(file, filename=attachment.filename)) + else: + large.append(attachment) + except HTTPException as e: + if e.status == 413: + large.append(attachment) + else: + raise + + if large: + embed = Embed(description=f"\n".join(f"[{attachment.filename}]({attachment.url})" for attachment in large)) + embed.set_footer(text="Attachments exceed upload size limit.") + await destination.send(embed=embed) diff --git a/config-default.yml b/config-default.yml index 0019d1688..1ecdfc5b9 100644 --- a/config-default.yml +++ b/config-default.yml @@ -308,11 +308,17 @@ reddit: subreddits: - 'r/Python' + wolfram: # Max requests per day. user_limit_day: 10 guild_limit_day: 67 key: !ENV "WOLFRAM_API_KEY" + +big_brother: + log_delay: 15 + + config: required_keys: ['bot.token'] |