yunke/yunke-server/app/scraper/media.py

"""Media downloader (M2, Task 2.6).

Fetches WeChat / OSS media blobs referenced by message records into the local
:class:`app.storage.object_store.ObjectStore` and records them in the
:class:`app.db.models.Media` table, deduplicated by the source ``md5``.

Phase-0 contract (see ``docs/superpowers/specs/phase0-findings.md`` §2/§3):

* WeChat type ``2`` (image / emoji, ``.webp``) and ``4`` (video, ``.mp4``) carry
  a DIRECT http(s) URL on ``file`` / ``fileTh`` / ``fileHd`` plus ``md5`` ->
  downloadable, deduped by ``md5``.
* WeChat type ``3`` (legacy image) carries an OPAQUE id on ``file`` (not a URL)
  -> cannot be resolved yet; :func:`download_message_media` returns ``None``.
* WeChat type ``10`` (system / link / card / invite) carries a
  ``support.weixin.qq.com`` URL that is a documentation link, NOT a media blob
  -> explicitly skipped before any network attempt.

The downloader is URL-driven and generic: any http(s) media URL works (the
recording URLs on ``yunke-pcfile`` are the same shape), so call recordings and
voice-call recordings can be fetched by the same path once their owning records
are wired in. That wiring is OUT of scope here (this module only consumes
normalized message dicts); see ``app.scraper.wechat`` for those.

Resilient by design: a download failure is logged and the function returns
``None`` rather than raising. A retry queue is a LATER milestone.
"""

from __future__ import annotations

import logging
import mimetypes
from pathlib import PurePosixPath
from typing import Any
from urllib.parse import urlparse

import httpx
import sqlalchemy as sa
from sqlalchemy.ext.asyncio import AsyncSession

from app.db.models import Media
from app.session.yunke_client import YunkeClient
from app.storage.object_store import ObjectStore
from app.storage.repos import upsert_media

__all__ = [
    "download_message_media",
    "download_url",
    "guess_mime_type",
    "media_type_from_message",
]

logger = logging.getLogger(__name__)

# Normalized message ``type`` label -> :attr:`Media.media_type` value. Only the
# Phase-0-confirmed image / video labels map today; voice / file codes
# (``6`` / ``9`` / ...) remain ``"unknown"`` at the scraper layer and yield
# ``None`` here until they are confirmed (phase0-findings.md §2).
_MEDIA_TYPE_BY_LABEL: dict[str, str] = {
    "image": "image",
    "image_legacy": "image",
    "video": "video",
}

# Extension -> MIME fallback map for the OSS / call-recording media types we
# observe. Used when the stdlib ``mimetypes`` database does not recognize the
# extension (e.g. ``.webp`` / ``.amr`` on some platforms) so content-type
# derivation does not depend on the system MIME database.
_EXT_MIME: dict[str, str] = {
    ".webp": "image/webp",
    ".jpg": "image/jpeg",
    ".jpeg": "image/jpeg",
    ".png": "image/png",
    ".gif": "image/gif",
    ".bmp": "image/bmp",
    ".mp4": "video/mp4",
    ".mp3": "audio/mpeg",
    ".m4a": "audio/mp4",
    ".wav": "audio/wav",
    ".amr": "audio/amr",
}


def media_type_from_message(msg: dict[str, Any]) -> str | None:
    """Map a normalized message dict to a :attr:`Media.media_type` label.

    Returns ``"image"`` for ``image`` / ``image_legacy`` messages and ``"video"``
    for ``video`` messages; ``None`` for everything else (text / system /
    unknown / not-yet-confirmed voice-or-file codes).
    """
    return _MEDIA_TYPE_BY_LABEL.get(msg.get("type"))  # type: ignore[arg-type]


def _is_http_url(value: Any) -> bool:
    """True iff ``value`` is an ``http://`` / ``https://`` string."""
    return isinstance(value, str) and value.startswith(("http://", "https://"))


def _pick_media_url(msg: dict[str, Any]) -> str | None:
    """Pick the best downloadable http(s) URL from a message dict.

    Preference order: ``file_url`` (original) -> ``file_hd_url`` (high-def) ->
    ``file_th_url`` (thumbnail). Returns ``None`` when no candidate is an
    http(s) URL (covers type-``3`` opaque ids and text messages).
    """
    for key in ("file_url", "file_hd_url", "file_th_url"):
        url = msg.get(key)
        if _is_http_url(url):
            return url  # type: ignore[return-value]
    return None


def _ext_from_url(url: str) -> str:
    """Extract the lowercase extension (with leading ``.``) from a URL path.

    Query / fragment are stripped first. Returns ``""`` when the path has no
    suffix (e.g. an OSS callback URL without an extension).
    """
    suffix = PurePosixPath(urlparse(url).path).suffix
    return suffix.lower()


def guess_mime_type(ext: str) -> str | None:
    """Guess a MIME type from a file extension.

    The built-in :data:`_EXT_MIME` map is consulted first (covers ``.webp`` /
    ``.mp4`` / ``.amr`` reliably regardless of platform), then the stdlib
    ``mimetypes`` database. Returns ``None`` when nothing matches.
    """
    if not ext:
        return None
    key = ext.lower()
    if key in _EXT_MIME:
        return _EXT_MIME[key]
    guessed, _ = mimetypes.guess_type(f"x{key}")
    return guessed


async def download_url(client: YunkeClient, url: str) -> bytes:
    """Fetch the raw bytes at ``url`` via the client's underlying httpx session.

    The :class:`YunkeClient` httpx client carries the session cookies (harmless
    for the public OSS media URLs, useful for any same-origin authenticated
    URLs) and is the transport ``respx`` mocks in the test suite, so media
    downloads go through the same mocked transport as every other scraper
    request. Redirects are followed (OSS / avatar URLs commonly redirect).
    """
    resp = await client.client.get(url, follow_redirects=True)
    resp.raise_for_status()
    return resp.content


async def download_message_media(
    client: YunkeClient,
    msg: dict[str, Any],
    store: ObjectStore,
    session: AsyncSession,
) -> str | None:
    """Download one message's media blob and record it in the ``Media`` table.

    Returns the dedup hash (the message ``md5`` when present, else the sha256 of
    the downloaded bytes), or ``None`` when:

    * the message is a type-``10`` system / link message (non-media; skipped
      before any network attempt), or
    * the message carries no downloadable http(s) URL (text, type-``3``
      opaque-id images, or any message whose ``file``/``fileHd``/``fileTh`` are
      absent or non-URL), or
    * the download fails (the error is logged; no exception escapes).

    Dedup: when the message carries an ``md5`` and a :class:`Media` row for that
    hash already exists with a local file, the download is SKIPPED and the hash
    is returned without touching the network.

    The blob is written via :meth:`ObjectStore.put_by_hash` (sha256
    content-addressed on disk), so identical bytes are stored once regardless of
    how many messages reference them.
    """
    # Type-10 system / link messages carry a support URL that is NOT media
    # (phase0-findings.md §2): skip them explicitly before URL selection so the
    # downloader never fetches a documentation page as if it were a media blob.
    if msg.get("type") == "system":
        return None

    url = _pick_media_url(msg)
    if url is None:
        return None

    md5 = msg.get("md5")
    if md5:
        existing = (
            await session.execute(
                sa.select(Media.local_path).where(Media.media_hash == md5)
            )
        ).scalar_one_or_none()
        if existing:
            # Already downloaded (local file present); short-circuit the network.
            return md5  # type: ignore[return-value]

    try:
        data = await download_url(client, url)
    except (httpx.HTTPError, OSError) as exc:
        logger.warning("media download failed for %s: %s", url, exc)
        return None

    ext = _ext_from_url(url)
    sha256_hex, path = await store.put_by_hash(data, ext=ext)
    media_hash: str = md5 if md5 else sha256_hex

    await upsert_media(
        session,
        media_hash=media_hash,
        channel=msg.get("channel") or "wechat",
        media_type=media_type_from_message(msg),
        original_url=msg.get("file_url"),
        original_th_url=msg.get("file_th_url"),
        original_hd_url=msg.get("file_hd_url"),
        local_path=str(path),
        size_bytes=len(data),
        mime_type=guess_mime_type(ext),
        raw_json=msg.get("raw_json"),
    )
    return media_hash