content-harvester/core/media.py

"""下载 m3u8/mp4 视频 + 封面图,上传到 MinIO。

MinIO 配置由调用方(主站经 lobster payload)传入,本模块不读全局配置。
敏感值(secret_key)不打印。
"""
from __future__ import annotations

import logging
import mimetypes
import os
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Optional
from urllib.parse import urlparse

logger = logging.getLogger("harvester.media")

_UA = (
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
    "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0 Safari/537.36"
)


# --------------------------------------------------------------------------
# MinIO 配置
# --------------------------------------------------------------------------

@dataclass
class MinioConfig:
    endpoint: str = ""        # 形如 https://minio.example.com
    access_key: str = ""
    secret_key: str = ""
    bucket: str = ""

    @classmethod
    def from_payload(cls, payload: dict) -> "MinioConfig":
        def g(*keys: str) -> str:
            for k in keys:
                v = payload.get(k)
                if v:
                    return str(v).strip()
            return ""
        cfg = cls(
            endpoint=g("minio_endpoint", "MINIO_ENDPOINT") or os.environ.get("MINIO_ENDPOINT", "").strip(),
            access_key=g("minio_access_key", "MINIO_ACCESS_KEY") or os.environ.get("MINIO_ACCESS_KEY", "").strip(),
            secret_key=g("minio_secret_key", "MINIO_SECRET_KEY") or os.environ.get("MINIO_SECRET_KEY", "").strip(),
            bucket=g("minio_bucket", "MINIO_BUCKET") or os.environ.get("MINIO_BUCKET", "").strip(),
        )
        return cfg

    @property
    def is_complete(self) -> bool:
        return bool(self.endpoint and self.access_key and self.secret_key and self.bucket)

    def safe_repr(self) -> str:
        return f"MinioConfig(endpoint={self.endpoint!r}, bucket={self.bucket!r}, access_key={self.access_key[:4]+'***' if self.access_key else ''!r}, secret_key=***)"


def _endpoint_host_secure(endpoint: str) -> tuple[str, bool]:
    """https://minio.example.com -> ('minio.aidfxn.com', True);裸 host 也兼容。"""
    e = endpoint.strip()
    if e.startswith("http://"):
        return e[len("http://"):].strip("/"), False
    if e.startswith("https://"):
        return e[len("https://"):].strip("/"), True
    # 没带 scheme,默认 https
    return e.strip("/"), True


def build_public_url(endpoint: str, bucket: str, object_key: str) -> str:
    return f"{endpoint.rstrip('/')}/{bucket.strip('/')}/{object_key.lstrip('/')}"


def upload_to_minio(local_path: Path, object_key: str, content_type: str, cfg: MinioConfig) -> str:
    """上传一个文件到 MinIO,返回公开 URL。失败抛异常。"""
    from minio import Minio  # 延迟 import,缺依赖时报错更清晰

    host, secure = _endpoint_host_secure(cfg.endpoint)
    client = Minio(host, access_key=cfg.access_key, secret_key=cfg.secret_key, secure=secure)
    client.fput_object(cfg.bucket, object_key, str(local_path), content_type=content_type or "application/octet-stream")
    return build_public_url(cfg.endpoint, cfg.bucket, object_key)


# --------------------------------------------------------------------------
# 下载
# --------------------------------------------------------------------------

_VIDEO_EXTS = (".mp4", ".mkv", ".webm", ".mov", ".ts", ".m4v")


def extract_video_id(detail_url: str) -> str:
    m = re.search(r"/(?:video|videos|photo|photos)/id-([A-Za-z0-9_-]+)(?:/\d+)?\.html", detail_url)
    if m:
        return m.group(1)
    # 兜底:用 url 的 hash
    import hashlib
    return "u" + hashlib.sha1(detail_url.encode("utf-8")).hexdigest()[:16]


# 兼容旧名;图集详情页用同一个 id 提取逻辑
extract_detail_id = extract_video_id


_IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".webp", ".gif", ".bmp")


def download_image(image_url: str, referer: str, dest_dir: Path, name: str) -> Optional[Path]:
    """流式下载一张图 → dest_dir/<name><ext>。失败返回 None(单张失败不致命)。"""
    import requests

    try:
        ext = os.path.splitext(urlparse(image_url).path)[1].lower()
        if ext not in _IMAGE_EXTS:
            ext = ".webp"
        dest_dir.mkdir(parents=True, exist_ok=True)
        path = dest_dir / f"{name}{ext}"
        with requests.get(image_url, headers={"Referer": referer, "User-Agent": _UA}, stream=True, timeout=30) as r:
            r.raise_for_status()
            with open(path, "wb") as f:
                for chunk in r.iter_content(chunk_size=64 * 1024):
                    if chunk:
                        f.write(chunk)
        if path.stat().st_size == 0:
            path.unlink(missing_ok=True)
            return None
        return path
    except Exception as exc:  # noqa: BLE001
        logger.warning("image download failed %s: %s", image_url, exc)
        return None


def head_ok(url: str, referer: str = "", timeout: int = 15) -> bool:
    """HEAD 探测一个 URL 是否可达(2xx)。失败/异常返回 False。"""
    import requests

    try:
        r = requests.head(url, headers={"Referer": referer or url, "User-Agent": _UA},
                          allow_redirects=True, timeout=timeout)
        if r.status_code == 405:  # 有些 CDN 不支持 HEAD,退而用 GET 流
            r = requests.get(url, headers={"Referer": referer or url, "User-Agent": _UA},
                             stream=True, timeout=timeout)
            r.close()
        return 200 <= r.status_code < 300
    except Exception:  # noqa: BLE001
        return False


def download_video(media_url: str, referer: str, dest_dir: Path, video_id: str,
                   concurrent_fragments: int = 8, max_mb: int = 0) -> Path:
    """用 yt-dlp 下载 m3u8/mp4 → dest_dir/<video_id>.mp4(尽量)。返回最终文件路径。

    max_mb > 0 时,下完发现超过则删文件并抛 ValueError("video too large: ...")。
    """
    import yt_dlp

    dest_dir.mkdir(parents=True, exist_ok=True)
    outtmpl = str(dest_dir / f"{video_id}.%(ext)s")
    opts = {
        "outtmpl": outtmpl,
        "format": "best",
        "merge_output_format": "mp4",
        "concurrent_fragment_downloads": max(1, int(concurrent_fragments)),
        "http_headers": {"Referer": referer, "User-Agent": _UA},
        "quiet": True,
        "no_warnings": True,
        "noprogress": True,
        "retries": 3,
        "fragment_retries": 3,
        "noplaylist": True,
    }
    with yt_dlp.YoutubeDL(opts) as ydl:
        ydl.download([media_url])

    # 找产物:dest_dir/<video_id>.<ext>
    candidates = sorted(
        (p for p in dest_dir.glob(f"{video_id}.*") if p.suffix.lower() in _VIDEO_EXTS),
        key=lambda p: p.stat().st_size,
        reverse=True,
    )
    if not candidates:
        raise RuntimeError("yt-dlp 未产出视频文件")
    path = candidates[0]
    size_mb = path.stat().st_size / (1024 * 1024)
    if max_mb and size_mb > max_mb:
        try:
            path.unlink()
        except OSError:
            pass
        raise ValueError(f"video too large: {size_mb:.0f}mb > {max_mb}mb")
    return path


def download_cover(cover_url: str, referer: str, dest_dir: Path, video_id: str) -> Optional[Path]:
    """流式下载封面图 → dest_dir/<video_id><ext>。失败返回 None(非致命)。"""
    import requests

    try:
        ext = os.path.splitext(urlparse(cover_url).path)[1].lower()
        if ext not in (".jpg", ".jpeg", ".png", ".webp", ".gif"):
            ext = ".jpg"
        dest_dir.mkdir(parents=True, exist_ok=True)
        path = dest_dir / f"{video_id}{ext}"
        with requests.get(cover_url, headers={"Referer": referer, "User-Agent": _UA}, stream=True, timeout=30) as r:
            r.raise_for_status()
            with open(path, "wb") as f:
                for chunk in r.iter_content(chunk_size=64 * 1024):
                    if chunk:
                        f.write(chunk)
        if path.stat().st_size == 0:
            path.unlink(missing_ok=True)
            return None
        return path
    except Exception as exc:  # noqa: BLE001
        logger.warning("cover download failed %s: %s", cover_url, exc)
        return None


def guess_content_type(path: Path, default: str = "application/octet-stream") -> str:
    return mimetypes.guess_type(str(path))[0] or default