Add discord file name normalization

1 files changed, 22 insertions, 0 deletions

diff --git a/bot/exts/utils/snekbox/_io.py b/bot/exts/utils/snekbox/_io.py
index 3ec5ff00a..88f10e4a2 100644
--- a/bot/exts/utils/snekbox/_io.py
+++ b/bot/exts/utils/snekbox/_io.py
@@ -6,6 +6,7 @@ from dataclasses import dataclass
 from io import BytesIO
 from pathlib import Path
 
+import regex
 from discord import File
 
 # Note discord bot upload limit is 8 MiB per file,
@@ -16,6 +17,14 @@ FILE_SIZE_LIMIT = 8 * 1024 * 1024
 FILE_COUNT_LIMIT = 10
 
 
+# ANSI escape sequences
+RE_ANSI = regex.compile(r"\x1b[^m]*m")
+# Characters with a leading backslash
+RE_BACKSLASH = regex.compile(r"\\.")
+# Discord disallowed file name characters
+RE_DISCORD_FILE_NAME_DISALLOWED = regex.compile(r"[^a-zA-Z0-9._-]+")
+
+
 def sizeof_fmt(num: int | float, suffix: str = "B") -> str:
     """Return a human-readable file size."""
     num = float(num)
@@ -28,6 +37,18 @@ def sizeof_fmt(num: int | float, suffix: str = "B") -> str:
     return f"{num_str} Yi{suffix}"
 
 
+def normalize_discord_file_name(name: str) -> str:
+    """Return a normalized valid discord file name."""
+    # Discord file names only allow A-Z, a-z, 0-9, underscores, dashes, and dots
+    # https://discord.com/developers/docs/reference#uploading-files
+    # Server will remove any other characters, but we'll get a 400 error for \ escaped chars
+    name = RE_ANSI.sub("_", name)
+    name = RE_BACKSLASH.sub("_", name)
+    # Replace any disallowed character with an underscore
+    name = RE_DISCORD_FILE_NAME_DISALLOWED.sub("_", name)
+    return name
+
+
 @dataclass
 class FileAttachment:
     """File Attachment from Snekbox eval."""
@@ -68,4 +89,5 @@ class FileAttachment:
     def to_file(self) -> File:
         """Convert to a discord.File."""
         name = Path(self.path).name
+        name = normalize_discord_file_name(name)
         return File(BytesIO(self.content), filename=name)