diff options
| -rw-r--r-- | bot/exts/utils/snekbox/_io.py | 22 |
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) |