Mirror live APstudent server messages

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)

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()