diff --git a/src/cogs/scam_image_detector.py b/src/cogs/scam_image_detector.py index b213a98..dc78c3e 100644 --- a/src/cogs/scam_image_detector.py +++ b/src/cogs/scam_image_detector.py @@ -6,7 +6,6 @@ import re import subprocess import tempfile -from collections import OrderedDict from dataclasses import dataclass from io import BytesIO from pathlib import Path @@ -63,8 +62,6 @@ def _env_ids(name: str, default: set) -> set: return parsed or set(default) -ALERT_USER_ID = _env_int("ALERT_USER_ID", 920819377627099166) -ALERT_MESSAGE = f"<@{ALERT_USER_ID}>" TARGET_APPLICATION_ID = 1508281890820460604 REVIEW_CHANNEL_ID = 1508284501485162541 MODERATION_BYPASS_ROLE_IDS = _env_ids( @@ -106,7 +103,6 @@ def has_moderation_bypass_role(member) -> bool: SCAN_CONCURRENCY = max(1, _env_int("SCAN_CONCURRENCY", 2)) SCAN_TIMEOUT_SECONDS = _env_int("SCAN_TIMEOUT_SECONDS", 60) -MAX_ALERTED_IDS = _env_int("MAX_ALERTED_IDS", 5000) DOWNLOAD_CHUNK_BYTES = 64 * 1024 HTTP_TIMEOUT_SECONDS = _env_int("HTTP_TIMEOUT_SECONDS", 20) @@ -1205,17 +1201,9 @@ def assess_media_bytes( class ScamImageDetector(commands.Cog): def __init__(self, bot: APBot) -> None: self.bot = bot - self.alerted_message_ids: "OrderedDict[int, None]" = OrderedDict() self._processing_ids: set = set() self._scan_semaphore = asyncio.Semaphore(SCAN_CONCURRENCY) - def _mark_alerted(self, message_id: Optional[int]) -> None: - if message_id is None: - return - self.alerted_message_ids[message_id] = None - while len(self.alerted_message_ids) > MAX_ALERTED_IDS: - self.alerted_message_ids.popitem(last=False) - def is_enabled_for_current_bot(self) -> bool: bot_user_id = getattr(getattr(self.bot, "user", None), "id", None) return bot_user_id == TARGET_APPLICATION_ID @@ -1461,10 +1449,10 @@ async def assess_message( return None, forwarded_sources return max(alert_assessments, key=lambda assessment: assessment.score), forwarded_sources - def build_alert_embed( + def build_review_embed( self, message: nextcord.Message, - assessment: ScamAssessment, + assessment: Optional[ScamAssessment], forwarded_sources: Iterable[ForwardedSource] = (), ) -> nextcord.Embed: original_content = getattr(message, "content", "") or "" @@ -1475,17 +1463,25 @@ def build_alert_embed( author_mention = getattr(author, "mention", str(author_id)) channel_mention = getattr(getattr(message, "channel", None), "mention", "unknown channel") + unsafe = assessment is not None + title = "Possible Unsafe Message Detected" if unsafe else "Server Message" + if is_forwarded: + title = f"Forwarded {title}" + embed = nextcord.Embed( - title="Possible Unsafe Forwarded Message Detected" - if is_forwarded - else "Possible Unsafe Message Detected", + title=title, description="The original message was not deleted.", - color=nextcord.Color.orange(), + color=nextcord.Color.orange() if unsafe else nextcord.Color.blue(), ) embed.add_field(name="User", value=f"{author_mention} (`{author_id}`)", inline=False) embed.add_field(name="Channel", value=channel_mention, inline=True) - embed.add_field(name="Score", value=f"{assessment.score}/100", inline=True) - embed.add_field(name="Reasons", value=clip("\n".join(f"- {reason}" for reason in assessment.reasons), 1000), inline=False) + if assessment is not None: + embed.add_field(name="Score", value=f"{assessment.score}/100", inline=True) + embed.add_field( + name="Reasons", + value=clip("\n".join(f"- {reason}" for reason in assessment.reasons), 1000), + inline=False, + ) if original_content: embed.add_field(name="Original Text", value=clip(original_content, 1000), inline=False) @@ -1497,7 +1493,7 @@ def build_alert_embed( if forwarded_text: embed.add_field(name="Forwarded Text", value=clip(forwarded_text, 1000), inline=False) - if assessment.scanned_text and assessment.content_kind != "text": + if assessment is not None and assessment.scanned_text and assessment.content_kind != "text": embed.add_field(name="Scanned Text", value=clip(f"`{assessment.scanned_text}`", 1000), inline=False) jump_url = getattr(message, "jump_url", None) @@ -1531,6 +1527,9 @@ async def _handle_message(self, message: nextcord.Message) -> None: if getattr(message, "guild", None) is None: return + if getattr(getattr(message, "channel", None), "id", None) == REVIEW_CHANNEL_ID: + return + author = getattr(message, "author", None) if getattr(author, "bot", False): return @@ -1541,10 +1540,11 @@ async def _handle_message(self, message: nextcord.Message) -> None: message_id = getattr(message, "id", None) if message_id is not None: - if message_id in self.alerted_message_ids or message_id in self._processing_ids: + if message_id in self._processing_ids: return self._processing_ids.add(message_id) + forwarded_sources = [] try: try: async with self._scan_semaphore: @@ -1553,7 +1553,7 @@ async def _handle_message(self, message: nextcord.Message) -> None: ) except asyncio.TimeoutError: log.warning("scan timed out for message %s", message_id) - return + assessment = None if assessment is None: return @@ -1565,11 +1565,10 @@ async def _handle_message(self, message: nextcord.Message) -> None: try: await review_channel.send( - content=ALERT_MESSAGE, - embed=self.build_alert_embed(message, assessment, forwarded_sources), - allowed_mentions=nextcord.AllowedMentions(users=True, roles=False, everyone=False), + content=None, + embed=self.build_review_embed(message, assessment, forwarded_sources), + allowed_mentions=nextcord.AllowedMentions.none(), ) - self._mark_alerted(message_id) except (nextcord.Forbidden, nextcord.HTTPException): pass finally: @@ -1581,17 +1580,13 @@ async def _handle_message(self, message: nextcord.Message) -> None: async def on_message(self, message: nextcord.Message) -> None: await self.handle_message(message) - @commands.Cog.listener() - async def on_message_edit(self, before: nextcord.Message, after: nextcord.Message) -> None: - await self.handle_message(after) - @commands.Cog.listener() async def on_raw_message_edit(self, payload: nextcord.RawMessageUpdateEvent) -> None: try: if not self.is_enabled_for_current_bot(): return - if payload.message_id in self.alerted_message_ids: + if payload.channel_id == REVIEW_CHANNEL_ID: return await asyncio.sleep(1) diff --git a/tests/test_scam_image_detector.py b/tests/test_scam_image_detector.py index a04f3b6..a44fd84 100644 --- a/tests/test_scam_image_detector.py +++ b/tests/test_scam_image_detector.py @@ -44,6 +44,7 @@ def make_message(channel_id=123, roles=None): embeds=[], reference=None, jump_url="https://discord.com/channels/1/123/789", + delete=AsyncMock(), ) @@ -135,7 +136,18 @@ def test_trusted_role_does_not_scan(): review_channel.send.assert_not_awaited() -def test_unsafe_message_from_any_channel_forwards_to_review_channel(): +def test_review_channel_messages_are_not_scanned_or_forwarded(): + bot, review_channel = make_bot() + cog = ScamImageDetector(bot) + cog.assess_message = AsyncMock(return_value=(dangerous_assessment(), [])) + + asyncio.run(cog.handle_message(make_message(channel_id=REVIEW_CHANNEL_ID))) + + cog.assess_message.assert_not_awaited() + review_channel.send.assert_not_awaited() + + +def test_unsafe_message_is_forwarded_without_ping_or_deletion(): bot, review_channel = make_bot() cog = ScamImageDetector(bot) cog.assess_message = AsyncMock(return_value=(dangerous_assessment(), [])) @@ -146,14 +158,19 @@ def test_unsafe_message_from_any_channel_forwards_to_review_channel(): cog.assess_message.assert_awaited_once_with(message) review_channel.send.assert_awaited_once() message.channel.send.assert_not_awaited() - assert review_channel.send.await_args.kwargs["content"] == "<@920819377627099166>" + message.delete.assert_not_awaited() + assert review_channel.send.await_args.kwargs["content"] is None + assert review_channel.send.await_args.kwargs["allowed_mentions"].users is False -def test_safe_message_is_not_forwarded(): +def test_safe_message_is_not_forwarded_or_deleted(): bot, review_channel = make_bot() cog = ScamImageDetector(bot) cog.assess_message = AsyncMock(return_value=(None, [])) + message = make_message(channel_id=999) - asyncio.run(cog.handle_message(make_message(channel_id=999))) + asyncio.run(cog.handle_message(message)) review_channel.send.assert_not_awaited() + message.channel.send.assert_not_awaited() + message.delete.assert_not_awaited()