scrapyard-cogworks/imagegrab/imagegrab.py

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