signal-bot/signal_bot/pack_ingest.py

# Copyright (c) 2026 Krzysztof Rudnicki. MIT licence.
"""Turning the attachments on one message into files inside the album.

Split out of pack_add.py at the 250-line cap, along the seam that was already
there: that half decides *whether* a command runs and what it answers, this
half does the reading, the stripping and the writing.

The order here is the load-bearing part. Media and manifests are written
first, the album is reloaded second, and only a reload that succeeds is
published. A reload that fails rolls all four files back and leaves the bot
serving the pack it was already serving -- and never raises, because
``Restart=always`` turns a startup-time ConfigError into a restart loop.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

from signal_bot.media_sniff import sniff
from signal_bot.pack_custom import write as write_custom
from signal_bot.pack_replies import ADDED, ALREADY_THERE, FAILED, NEW_TAG_HINT
from signal_bot.pack_store import (
    content_name,
    record_adds,
    restore,
    snapshot,
    write_media,
    write_tags,
)

if TYPE_CHECKING:
    from signal_bot.media_strip import Stripper
    from signal_bot.pack_add import AddResolver
    from signal_bot.pack_holder import PackHolder
    from signal_bot.registry import CommandContext

logger = logging.getLogger(__name__)


async def ingest(
    ctx: CommandContext, holder: PackHolder, stripper: Stripper
) -> tuple[dict[str, bytes], int]:
    """Fetch, sniff, strip and name every attachment on the message.

    Returns:
        The files that are new to the album, and how many were already in it.
    """
    fresh: dict[str, bytes] = {}
    duplicates = 0
    known = set(holder.pack.tags)
    for attachment_id in ctx.message.attachment_ids:
        payload = await _fetch(ctx, attachment_id)
        cleaned = _clean(payload, stripper)
        if cleaned is not None:
            name, body = cleaned
            if name in known or name in fresh:
                duplicates += 1
            else:
                fresh[name] = body
    return fresh, duplicates


async def _fetch(ctx: CommandContext, attachment_id: str) -> bytes | None:
    """Read one attachment, then delete signal-cli's own copy of it.

    Deleted whether the upload is accepted or refused: the point is that no
    copy survives outside the album, and a refused file is exactly the one
    nobody wants left lying in the container.
    """
    try:
        payload = await ctx.transport.fetch_attachment(attachment_id)
    except Exception:
        logger.exception("could not fetch an attachment")
        payload = None
    try:
        await ctx.transport.delete_attachment(attachment_id)
    except Exception:
        logger.exception("could not delete the server-side attachment copy")
    return payload


def _clean(payload: bytes | None, stripper: Stripper) -> tuple[str, bytes] | None:
    """Sniff and strip one upload, or None when it is not storable."""
    if payload is None:
        return None
    kind = sniff(payload)
    if kind is None:
        logger.warning("refusing an upload whose format is not on the allowlist")
        return None
    cleaned = stripper.strip(payload, kind)
    if cleaned is None:
        return None
    return content_name(cleaned, kind.suffix), cleaned


async def commit(
    ctx: CommandContext,
    holder: PackHolder,
    resolver: AddResolver,
    keys: tuple[str, ...],
    fresh: dict[str, bytes],
    duplicates: int,
) -> None:
    """Write the files and manifests, swap the album, or roll it all back."""
    pack = holder.pack
    directory = pack.directory
    previous = snapshot(directory)
    tags = {name: frozenset(existing) for name, existing in pack.tags.items()}
    for name, body in fresh.items():
        write_media(directory, name, body)
        tags[name] = frozenset(keys)
    write_tags(directory, tags)
    write_custom(directory, resolver.committed(pack.custom))
    record_adds(
        directory,
        {name: {"by": ctx.message.source_uuid, "at": int(ctx.now)} for name in fresh},
    )
    if not holder.reload():
        restore(directory, previous, tuple(fresh))
        await ctx.reply(FAILED)
        return
    await ctx.reply(_receipt(holder, keys, duplicates, minted=bool(resolver.pending)))


def _receipt(holder: PackHolder, keys: tuple[str, ...], duplicates: int, *, minted: bool) -> str:
    """Counts and pseudonyms only -- never a label, never a filename."""
    pack = holder.pack
    parts = [ADDED]
    parts.extend(f"{key}: {len(pack.matching(frozenset({key})))}" for key in keys)
    if duplicates:
        parts.append(ALREADY_THERE.format(count=duplicates))
    line = " · ".join(parts)
    return f"{line}\n{NEW_TAG_HINT}" if minted else line