From 1ee79a957699dd0a4af10a4b47acd8b91f63f030 Mon Sep 17 00:00:00 2001
From: Codex <codex@openai.com>
Date: Sat, 7 Jun 2025 00:13:50 +0000
Subject: [PATCH] feat(aimod): process direct link attachments

---
 cogs/aimod_cog.py | 80 +++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 78 insertions(+), 2 deletions(-)

diff --git a/cogs/aimod_cog.py b/cogs/aimod_cog.py
index 6b07a4f..b66a40b 100644
--- a/cogs/aimod_cog.py
+++ b/cogs/aimod_cog.py
@@ -18,6 +18,7 @@ import shutil  # For backing up files
 from typing import Optional, List, Dict, Any, Tuple  # For type hinting
 import asyncio
 import aiofiles
+import re
 
 # Google Generative AI Imports (using Vertex AI backend)
 from google import genai
@@ -446,6 +447,55 @@ class AIModerationCog(commands.Cog):
             print(f"Error processing video: {e}")
             return None, None
 
+    async def process_url_attachment(self, url: str) -> tuple[str, bytes, str, str]:
+        """Fetch an attachment from a direct link."""
+        import aiohttp
+
+        try:
+            cleaned_url = url.strip("<>")
+            filename = cleaned_url.split("/")[-1].split("?")[0]
+            _, ext = os.path.splitext(filename.lower())
+            if ext in self.image_extensions:
+                attachment_type = "image"
+            elif ext in self.gif_extensions:
+                attachment_type = "gif"
+            elif ext in self.video_extensions:
+                attachment_type = "video"
+            else:
+                return None, None, None, None
+
+            async with aiohttp.ClientSession() as session:
+                async with session.get(cleaned_url) as resp:
+                    if resp.status != 200:
+                        print(
+                            f"Failed to fetch URL attachment {cleaned_url}: {resp.status}"
+                        )
+                        return None, None, None, None
+                    data = await resp.read()
+                    mime_type = resp.headers.get(
+                        "Content-Type", f"image/{ext.lstrip('.')}"
+                    )
+                    return mime_type, data, attachment_type, filename
+        except Exception as e:
+            print(f"Error processing URL attachment {url}: {e}")
+            return None, None, None, None
+
+    def extract_direct_attachment_urls(self, text: str) -> List[str]:
+        """Return a list of direct image/video URLs found in the text."""
+
+        urls = re.findall(r"https?://\S+", text or "")
+        allowed_exts = (
+            self.image_extensions + self.gif_extensions + self.video_extensions
+        )
+        results = []
+        for u in urls:
+            cleaned = u.strip("<>")
+            path = cleaned.split("?")[0]
+            _, ext = os.path.splitext(path.lower())
+            if ext in allowed_exts:
+                results.append(cleaned)
+        return results
+
     # --- AI Moderation Command Group ---
     aimod_group = app_commands.Group(
         name="aimod", description="AI Moderation commands."
@@ -2077,8 +2127,13 @@ CRITICAL: Do NOT output anything other than the required JSON response.
         if message.author.bot:
             print(f"Ignoring message {message.id} from bot.")
             return
-        # Ignore messages without content or attachments
-        if not message.content and not message.attachments:
+        link_urls = (
+            self.extract_direct_attachment_urls(message.content)
+            if message.content
+            else []
+        )
+        # Ignore messages without content, attachments, or direct attachment links
+        if not message.content and not message.attachments and not link_urls:
             print(f"Ignoring message {message.id} with no content or attachments.")
             return
         # Ignore DMs
@@ -2125,6 +2180,27 @@ CRITICAL: Do NOT output anything other than the required JSON response.
                     f"Processed {len(image_data_list)} attachments for message {message.id}"
                 )
 
+        # Check for direct link attachments in the message content
+        if link_urls:
+            processed_links = 0
+            for url in link_urls:
+                mime_type, image_bytes, attachment_type, filename = (
+                    await self.process_url_attachment(url)
+                )
+                if mime_type and image_bytes and attachment_type:
+                    image_data_list.append(
+                        (mime_type, image_bytes, attachment_type, filename)
+                    )
+                    processed_links += 1
+                    print(
+                        f"Processed linked attachment: {filename} as {attachment_type}"
+                    )
+
+            if processed_links > 0:
+                print(
+                    f"Processed {processed_links} linked attachments for message {message.id}"
+                )
+
         # Only proceed with AI analysis if there's text to analyze or attachments
         if not message_content and not image_data_list:
             print(