diff --git a/imagegrab/__init__.py b/imagegrab/__init__.py new file mode 100644 index 0000000..146097e --- /dev/null +++ b/imagegrab/__init__.py @@ -0,0 +1,5 @@ +from .imagegrab import ImageGrab + + +async def setup(bot): + await bot.add_cog(ImageGrab(bot)) diff --git a/imagegrab/imagegrab.py b/imagegrab/imagegrab.py new file mode 100644 index 0000000..22c547a --- /dev/null +++ b/imagegrab/imagegrab.py @@ -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 diff --git a/imagegrab/info.json b/imagegrab/info.json new file mode 100644 index 0000000..b9dafc1 --- /dev/null +++ b/imagegrab/info.json @@ -0,0 +1,11 @@ +{ + "author": ["Scrapyard Cogworks"], + "install_msg": "ImageGrab cog installed. Use `[p]imagegrab` to zip all images in a channel.", + "name": "ImageGrab", + "short": "Zip all images from a channel with type-based filenames.", + "description": "Pulls every image attachment from a channel's history and sends them back as a zip file. If a message has text, images are named after that text (used as type tags for WebTable card decks).", + "tags": ["images", "zip", "download", "channel", "webtable", "cards"], + "requirements": ["aiohttp"], + "min_bot_version": "3.5.0", + "end_user_data_statement": "This cog does not store any end user data." +}