signal-bot/signal_bot/media_sniff.py

# Copyright (c) 2026 Krzysztof Rudnicki. MIT licence.
"""What an uploaded file actually is, decided from its bytes.

The filename a phone sends is not evidence: it is attacker-chosen, it is
discarded on the way in (the stored name is a content hash), and ".jpg" on a
file that is not a JPEG would be stored, served, and then refused by
``pack_load`` on the *next* start -- taking the bot down for a photo that
looked fine when it arrived. So the suffix is derived here instead, from the
magic bytes, and anything not on the allowlist never reaches the album.

The ISO base-media formats (mp4, mov, m4v, heic, heif, avif) share one
container and are told apart by the brand in the ``ftyp`` box rather than by
extension, which is the only thing that distinguishes an iPhone photo from an
iPhone video at the byte level.

SVG is the one member with no magic number at all -- it is XML -- so it is
matched on a leading prefix instead, and it is the one format where that is
the whole test.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Final

# How much of a file the sniffer ever looks at. The longest test is the
# ISO-BMFF brand, which ends at byte 12; the SVG prefix scan wants a little
# room for leading whitespace and a byte-order mark.
HEADER_BYTES: Final = 64

# A byte-order mark, which a phone's XML writer may emit ahead of "<?xml".
_BOM: Final = b"\xef\xbb\xbf"

_SVG_PREFIXES: Final = (b"<?xml", b"<svg")


@dataclass(frozen=True, slots=True)
class MediaKind:
    """One accepted format: how it is stored and how it is served."""

    # The suffix the stored file gets, always lowercase and always with a dot.
    suffix: str
    # What Signal is told the file is when it is posted back.
    media_type: str
    # Whether exiftool 13.55 can write this container at all. False means the
    # file is stored exactly as it arrived -- see media_strip.strip_bytes and
    # the warning in DOCS-add-photo.md.
    strippable: bool = True


JPEG: Final = MediaKind(".jpg", "image/jpeg")
PNG: Final = MediaKind(".png", "image/png")
GIF: Final = MediaKind(".gif", "image/gif")
WEBP: Final = MediaKind(".webp", "image/webp")
MP4: Final = MediaKind(".mp4", "video/mp4")
MOV: Final = MediaKind(".mov", "video/quicktime")
M4V: Final = MediaKind(".m4v", "video/x-m4v")
HEIC: Final = MediaKind(".heic", "image/heic")
HEIF: Final = MediaKind(".heif", "image/heif")
AVIF: Final = MediaKind(".avif", "image/avif")
# The two exiftool 13.55 cannot write. Kept on the allowlist by kuhy's
# decision of 2026-08-31; they are stored with their metadata intact.
WEBM: Final = MediaKind(".webm", "video/webm", strippable=False)
SVG: Final = MediaKind(".svg", "image/svg+xml", strippable=False)

# Every kind the album accepts, which is also every kind pack.MEDIA_TYPES has
# to know about: a file stored here and not there is refused by the *next*
# load, long after the upload that caused it.
KINDS: Final[tuple[MediaKind, ...]] = (
    JPEG,
    PNG,
    GIF,
    WEBP,
    MP4,
    MOV,
    M4V,
    HEIC,
    HEIF,
    AVIF,
    WEBM,
    SVG,
)

# Leading-byte signatures, longest-first so no prefix shadows another.
_SIGNATURES: Final[tuple[tuple[bytes, MediaKind], ...]] = (
    (b"\x89PNG\r\n\x1a\n", PNG),
    (b"\xff\xd8\xff", JPEG),
    (b"GIF87a", GIF),
    (b"GIF89a", GIF),
    (b"\x1a\x45\xdf\xa3", WEBM),
)

# The ``ftyp`` brand at bytes 8..12, lowercased. Everything ISO-BMFF that is
# not listed is treated as plain mp4, which is what the brand table's own
# "iso*"/"mp4*" families all are.
_BRANDS: Final[dict[bytes, MediaKind]] = {
    b"heic": HEIC,
    b"heix": HEIC,
    b"hevc": HEIC,
    b"hevx": HEIC,
    b"heim": HEIC,
    b"heis": HEIC,
    b"mif1": HEIF,
    b"msf1": HEIF,
    b"heif": HEIF,
    b"avif": AVIF,
    b"avis": AVIF,
    b"qt  ": MOV,
    b"m4v ": M4V,
}


def sniff(payload: bytes) -> MediaKind | None:
    """Identify one upload, or None when it is not a format we accept.

    Args:
        payload: The file exactly as it arrived.

    Returns:
        The matching kind, or None for anything off the allowlist.
    """
    header = payload[:HEADER_BYTES]
    for signature, kind in _SIGNATURES:
        if header.startswith(signature):
            return kind
    if _is_riff_webp(header):
        return WEBP
    brand = _iso_brand(header)
    if brand is not None:
        return _BRANDS.get(brand, MP4)
    return SVG if _is_svg(header) else None


def _is_riff_webp(header: bytes) -> bool:
    """True for a RIFF container whose form type is WEBP.

    Both halves are checked: RIFF also carries wav and avi, and a bare
    "RIFF" test would file either of them as a picture.
    """
    return header.startswith(b"RIFF") and header[8:12] == b"WEBP"


def _iso_brand(header: bytes) -> bytes | None:
    """The lowercased ISO base-media brand, or None if this is not one.

    The box length at bytes 0..4 is deliberately not validated. It varies
    between writers and is not what identifies the format; "ftyp" at byte 4
    is, and every file in this family has it.
    """
    if header[4:8] != b"ftyp":
        return None
    return header[8:12].lower()


def _is_svg(header: bytes) -> bool:
    """True for something that opens like an SVG document.

    A prefix test rather than a parse: this decides *which* allowlist entry a
    file belongs to, not whether it is well-formed, and a document that opens
    with "<?xml" and turns out to be some other XML is stored as an .svg that
    Signal simply will not render. That is a bad photo, not a leak.
    """
    text = header.removeprefix(_BOM).lstrip()
    return text.startswith(_SVG_PREFIXES)