signal-bot/signal_bot/pack_store.py

# Copyright (c) 2026 Krzysztof Rudnicki. MIT licence.
"""Putting a file into the album, and taking it back out when the swap fails.

Every write here is temp-then-``Path.replace``, the pattern
``pack_album.Album.write`` already uses, because a manifest that is half
written is a manifest that fails the next load -- and with ``Restart=always``
a startup-time ConfigError is a loop, not an error message.

The stored name is a content hash. The original filename is discarded: it is
attacker-chosen, it is frequently somebody's name and a date, and hashing
makes "is this photo already here?" a lookup instead of a comparison.

``adds.json`` records who added what. Nothing reads it -- ``pack_load``
ignores the file entirely -- and it exists because ``!undo_photo`` cannot be
built later without it, and the id has to have been recorded at add time.
It stays inside the album directory, which is 0700 and outside the
repository, so it is subject to exactly the same handling as the photos.
"""

from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Final

from signal_bot.pack import MEDIA_DIR, TAGS_FILE
from signal_bot.pack_custom import CUSTOM_FILE, FILE_MODE

if TYPE_CHECKING:
    from collections.abc import Mapping
    from pathlib import Path

ADDS_FILE: Final = "adds.json"

# Half a SHA-256 is 128 bits of name. Collisions are not a threat model here;
# a shorter name is simply easier to read in a directory listing.
NAME_LENGTH: Final = 32


def content_name(payload: bytes, suffix: str) -> str:
    """The filename one upload gets, derived from the bytes themselves."""
    digest = hashlib.sha256(payload).hexdigest()[:NAME_LENGTH]
    return f"{digest}{suffix}"


@dataclass(frozen=True, slots=True)
class Snapshot:
    """The album's manifests as they were before a write, for rolling back.

    Bytes rather than parsed documents: restoring has to put back exactly
    what was there, and a reparse-and-redump could quietly normalise a file
    the tagger wrote by hand.
    """

    tags: bytes | None
    custom: bytes | None
    adds: bytes | None


def snapshot(directory: Path) -> Snapshot:
    """Read the three manifests as they stand, before anything is written."""
    return Snapshot(
        tags=_read(directory / TAGS_FILE),
        custom=_read(directory / CUSTOM_FILE),
        adds=_read(directory / ADDS_FILE),
    )


def restore(directory: Path, previous: Snapshot, written: tuple[str, ...]) -> None:
    """Undo a failed add: put the manifests back and delete the new files.

    Called only when the album no longer loads after a write, which means the
    bot is still serving the *old* pack from memory. Leaving the new files
    behind would make the next start fail on a tag that no longer resolves.
    """
    _restore(directory / TAGS_FILE, previous.tags)
    _restore(directory / CUSTOM_FILE, previous.custom)
    _restore(directory / ADDS_FILE, previous.adds)
    for name in written:
        (directory / MEDIA_DIR / name).unlink(missing_ok=True)


def write_media(directory: Path, name: str, payload: bytes) -> None:
    """Put one stripped file into the album, atomically."""
    album = directory / MEDIA_DIR
    album.mkdir(parents=True, exist_ok=True)
    path = album / name
    temporary = path.with_name(f"{name}.tmp")
    temporary.write_bytes(payload)
    temporary.chmod(FILE_MODE)
    temporary.replace(path)


def write_tags(directory: Path, tags: Mapping[str, frozenset[str]]) -> None:
    """Replace tags.json with the merged table."""
    document = {name: sorted(keys) for name, keys in sorted(tags.items())}
    write_json(directory / TAGS_FILE, document)


def record_adds(directory: Path, entries: Mapping[str, dict[str, Any]]) -> None:
    """Append to adds.json, which nothing at run time ever reads back."""
    document = _document(directory / ADDS_FILE)
    document.update(entries)
    write_json(directory / ADDS_FILE, document)


def write_json(path: Path, document: Mapping[str, object]) -> None:
    """Write one manifest atomically, 0600, so a crash cannot halve it."""
    temporary = path.with_name(f"{path.name}.tmp")
    temporary.write_text(json.dumps(document, indent=2), encoding="utf-8")
    temporary.chmod(FILE_MODE)
    temporary.replace(path)


def _document(path: Path) -> dict[str, Any]:
    """One JSON object, or an empty one for a file that is not there yet.

    Deliberately forgiving, unlike pack_load: adds.json is write-only
    bookkeeping, and refusing to store a photo because an unread audit file
    got corrupted would be the tail wagging the dog.
    """
    raw = _read(path)
    if raw is None:
        return {}
    try:
        document = json.loads(raw)
    except json.JSONDecodeError:
        return {}
    return document if isinstance(document, dict) else {}


def _read(path: Path) -> bytes | None:
    """The file's bytes, or None when it does not exist."""
    try:
        return path.read_bytes()
    except OSError:
        return None


def _restore(path: Path, previous: bytes | None) -> None:
    """Put one manifest back, or remove it if there was none."""
    if previous is None:
        path.unlink(missing_ok=True)
        return
    path.write_bytes(previous)