Add imagegrab cog: zip channel images with type-based filenames
Images are named using the message text as the filename (type tags for WebTable card decks). Single image = text.ext, multiple = text_1.ext, text_2.ext, etc. Falls back to messageID.ext when no text is present.
This commit is contained in:
@@ -0,0 +1,187 @@
|
||||
import io
|
||||
import os
|
||||
import zipfile
|
||||
from typing import Optional
|
||||
|
||||
import aiohttp
|
||||
import discord
|
||||
from redbot.core import commands, checks
|
||||
|
||||
|
||||
class ImageGrab(commands.Cog):
|
||||
"""Grab all images from a channel and bundle them into a zip.
|
||||
|
||||
If a message contains text alongside images, that text becomes the
|
||||
filename (used as type tags for WebTable card decks).
|
||||
Example: message text "gaming, peeboo" with 3 images produces
|
||||
gaming, peeboo_1.png, gaming, peeboo_2.png, gaming, peeboo_3.png
|
||||
"""
|
||||
|
||||
def __init__(self, bot):
|
||||
self.bot = bot
|
||||
self.session = aiohttp.ClientSession()
|
||||
|
||||
async def cog_unload(self):
|
||||
await self.session.close()
|
||||
|
||||
@commands.command()
|
||||
@checks.mod_or_permissions(manage_messages=True)
|
||||
@commands.bot_has_permissions(attach_files=True, read_message_history=True)
|
||||
async def imagegrab(
|
||||
self,
|
||||
ctx: commands.Context,
|
||||
channel: Optional[discord.TextChannel] = None,
|
||||
limit: Optional[int] = None,
|
||||
):
|
||||
"""Zip all images from a channel.
|
||||
|
||||
If a message has text, images are named after it (type tags).
|
||||
Multiple images in one message get _1, _2, etc. suffixes.
|
||||
|
||||
`channel` defaults to the current channel.
|
||||
`limit` caps how many messages to scan (omit for entire history).
|
||||
"""
|
||||
channel = channel or ctx.channel
|
||||
|
||||
status = await ctx.send(f"Scanning **#{channel.name}** for images… this may take a while.")
|
||||
|
||||
image_exts = (".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp", ".tiff")
|
||||
images = [] # list of (filename, url)
|
||||
|
||||
count = 0
|
||||
async for message in channel.history(limit=limit, oldest_first=True):
|
||||
# Collect all image attachments from this message
|
||||
msg_images = []
|
||||
for att in message.attachments:
|
||||
if att.filename.lower().endswith(image_exts):
|
||||
ext = os.path.splitext(att.filename)[1].lower()
|
||||
msg_images.append((ext, att.url))
|
||||
|
||||
# Also grab image embeds (linked images)
|
||||
for embed in message.embeds:
|
||||
if embed.type == "image" and embed.url:
|
||||
raw_ext = embed.url.split("?")[0].rsplit(".", 1)[-1].lower()
|
||||
if f".{raw_ext}" in image_exts:
|
||||
msg_images.append((f".{raw_ext}", embed.url))
|
||||
|
||||
if not msg_images:
|
||||
count += 1
|
||||
if count % 2000 == 0:
|
||||
await status.edit(content=f"Scanned {count} messages, found {len(images)} images so far…")
|
||||
continue
|
||||
|
||||
# Determine filename base from message text
|
||||
msg_text = message.content.strip() if message.content else ""
|
||||
|
||||
if msg_text:
|
||||
# Use message text as filename base
|
||||
# Sanitize for filesystem but keep commas/spaces (they're the type info)
|
||||
base = _sanitize_filename(msg_text)
|
||||
|
||||
if len(msg_images) == 1:
|
||||
ext, url = msg_images[0]
|
||||
images.append((f"{base}{ext}", url))
|
||||
else:
|
||||
for i, (ext, url) in enumerate(msg_images, 1):
|
||||
images.append((f"{base}_{i}{ext}", url))
|
||||
else:
|
||||
# No text — use original filename prefixed with message ID to avoid collisions
|
||||
for ext, url in msg_images:
|
||||
# For attachments we can get the original name from the URL
|
||||
images.append((f"{message.id}{ext}", url))
|
||||
|
||||
count += 1
|
||||
if count % 2000 == 0:
|
||||
await status.edit(content=f"Scanned {count} messages, found {len(images)} images so far…")
|
||||
|
||||
if not images:
|
||||
await status.edit(content="No images found in that channel.")
|
||||
return
|
||||
|
||||
await status.edit(content=f"Found {len(images)} images. Downloading and zipping…")
|
||||
|
||||
# Download all images
|
||||
downloaded = [] # list of (filename, bytes)
|
||||
failed = 0
|
||||
for filename, url in images:
|
||||
try:
|
||||
async with self.session.get(url) as resp:
|
||||
if resp.status == 200:
|
||||
downloaded.append((filename, await resp.read()))
|
||||
else:
|
||||
failed += 1
|
||||
except Exception:
|
||||
failed += 1
|
||||
|
||||
if not downloaded:
|
||||
await status.edit(content="All image downloads failed.")
|
||||
return
|
||||
|
||||
# Build zip chunks that stay under the upload limit.
|
||||
# Use 9 MB as the target to leave headroom for zip overhead.
|
||||
max_chunk = 9 * 1024 * 1024
|
||||
|
||||
def _finalize(zf, buf):
|
||||
"""Close the zip and return the seeked buffer."""
|
||||
zf.close()
|
||||
buf.seek(0)
|
||||
return buf
|
||||
|
||||
def _new_zip():
|
||||
"""Create a fresh BytesIO + ZipFile pair."""
|
||||
buf = io.BytesIO()
|
||||
return zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED), buf
|
||||
|
||||
chunks: list[io.BytesIO] = []
|
||||
current_zf, current_buf = _new_zip()
|
||||
current_size = 0
|
||||
|
||||
for filename, data in downloaded:
|
||||
file_size = len(data)
|
||||
|
||||
# If adding this file would exceed the limit, finalize current chunk first
|
||||
if current_size > 0 and current_size + file_size > max_chunk:
|
||||
chunks.append(_finalize(current_zf, current_buf))
|
||||
current_zf, current_buf = _new_zip()
|
||||
current_size = 0
|
||||
|
||||
current_zf.writestr(filename, data)
|
||||
current_size += file_size
|
||||
|
||||
# Finalize the last chunk
|
||||
if current_size > 0:
|
||||
chunks.append(_finalize(current_zf, current_buf))
|
||||
else:
|
||||
current_zf.close()
|
||||
|
||||
result = f"Here are **{len(downloaded)}** images from **#{channel.name}** in **{len(chunks)}** zip(s)."
|
||||
if failed:
|
||||
result += f" ({failed} failed to download)"
|
||||
await status.edit(content=result)
|
||||
|
||||
for i, buf in enumerate(chunks, 1):
|
||||
suffix = f"_part{i}" if len(chunks) > 1 else ""
|
||||
await ctx.send(
|
||||
content=f"Part {i}/{len(chunks)}" if len(chunks) > 1 else None,
|
||||
file=discord.File(buf, filename=f"{channel.name}_images{suffix}.zip"),
|
||||
)
|
||||
|
||||
|
||||
def _sanitize_filename(text: str) -> str:
|
||||
"""Sanitize text for use as a filename while preserving commas and spaces.
|
||||
|
||||
Removes/replaces characters that are invalid in filenames on most OS.
|
||||
Keeps commas, spaces, hyphens, and underscores.
|
||||
"""
|
||||
# Remove characters that are problematic in filenames
|
||||
invalid = '<>:"/\\|?*\x00'
|
||||
result = ""
|
||||
for ch in text:
|
||||
if ch in invalid:
|
||||
continue
|
||||
result += ch
|
||||
# Collapse multiple spaces
|
||||
result = " ".join(result.split())
|
||||
# Trim to reasonable length (255 is typical FS limit, leave room for suffix + ext)
|
||||
result = result[:200]
|
||||
return result
|
||||
Reference in New Issue
Block a user