content-harvester/server/app.py

"""content-harvester —— 内容采集服务,实现主站 lobster 执行器的调用契约。

当前内置一个 xchina.co 的 scraper(core/scraper.py),以后可加更多来源。
对接 app/services/lobster_client.py::call_lobster_agent:
  POST /agents/{agent_id}/invoke   body = 主站合并后的 payload, 返回 {"task_id": ...}
  GET  /tasks/{task_id}            返回 {"status": running|succeeded|failed, "output": ..., "error": ...}

鉴权: 主站带 `Authorization: Bearer <lobster_api_key>`,需与本服务 env HARVESTER_TOKEN 一致。
SSRF 防护: 只接受 ALLOWED_HOSTS(默认 xchina.co 及其子域) 的 content_url。
"""
from __future__ import annotations

import logging
import os
from urllib.parse import urlparse

from fastapi import FastAPI, Header, HTTPException
from pydantic import BaseModel

from core.media import MinioConfig
from server import jobs

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s")
logger = logging.getLogger("harvester.app")

SERVICE_TOKEN = os.environ.get("HARVESTER_TOKEN", "")
HEADLESS = os.environ.get("HEADLESS", "1") not in ("0", "false", "False")
ALLOWED_HOSTS = tuple(
    h.strip().lower() for h in os.environ.get("ALLOWED_HOSTS", "xchina.co").split(",") if h.strip()
)
MAX_PAGES = int(os.environ.get("MAX_PAGES", "50"))

app = FastAPI(title="content-harvester", version="1.0")


def _check_auth(authorization: str | None) -> None:
    if not SERVICE_TOKEN:
        # 未配置 token 时不鉴权(仅适合内网/本地);生产务必设 HARVESTER_TOKEN
        return
    expected = f"Bearer {SERVICE_TOKEN}"
    if authorization != expected:
        raise HTTPException(status_code=401, detail="invalid token")


def _host_allowed(url: str) -> bool:
    try:
        host = (urlparse(url).hostname or "").lower()
    except Exception:
        return False
    if not host:
        return False
    return any(host == h or host.endswith("." + h) for h in ALLOWED_HOSTS)


class TaskStatusResponse(BaseModel):
    status: str
    output: dict | None = None
    error: str | None = None


@app.get("/healthz")
async def healthz():
    return {"ok": True, "headless": HEADLESS, "allowed_hosts": list(ALLOWED_HOSTS)}


@app.post("/agents/{agent_id}/invoke")
async def invoke(agent_id: str, payload: dict, authorization: str | None = Header(default=None)):
    _check_auth(authorization)

    # payload 即主站 task.input_params:含 content_url / platform_name / source_kind / start_page / end_page / _crawl_config ...
    source_kind = payload.get("source_kind")
    if source_kind is not None and str(source_kind) != "website":
        raise HTTPException(status_code=400, detail=f"unsupported source_kind: {source_kind}(本服务只处理 website)")

    content_url = (payload.get("content_url") or payload.get("series_url") or payload.get("url") or "").strip()
    if not content_url:
        raise HTTPException(status_code=400, detail="missing content_url")
    if not _host_allowed(content_url):
        raise HTTPException(status_code=400, detail=f"host not allowed: {content_url}")

    def _as_int(v, default):
        try:
            return int(v)
        except (TypeError, ValueError):
            return default

    start_page = max(1, _as_int(payload.get("start_page"), 1))
    end_page = max(start_page, _as_int(payload.get("end_page"), start_page))
    if end_page - start_page + 1 > MAX_PAGES:
        end_page = start_page + MAX_PAGES - 1

    source_name = str(payload.get("platform_name") or "xchina")[:100]

    # MinIO 配置由主站经 lobster payload 下发(主站把绑定的 minio_* Skill 变量并入 payload);
    # 缺失则降级:不上传,video_url 保留原始 m3u8 地址
    minio_cfg = MinioConfig.from_payload(payload)

    job = jobs.submit(content_url, source_name, start_page, end_page, headless=HEADLESS, minio_cfg=minio_cfg)
    logger.info(
        "invoke agent=%s job=%s url=%s pages=%s-%s minio=%s",
        agent_id, job.id, content_url, start_page, end_page,
        "on" if minio_cfg.is_complete else "off",
    )
    return {"task_id": job.id, "status": job.status}


@app.get("/tasks/{task_id}")
async def get_task(task_id: str, authorization: str | None = Header(default=None)):
    _check_auth(authorization)
    job = jobs.get_job(task_id)
    if not job:
        raise HTTPException(status_code=404, detail="task not found")
    return job.to_dict()


@app.post("/tasks/{task_id}/cancel")
async def cancel_task(task_id: str, authorization: str | None = Header(default=None)):
    _check_auth(authorization)
    ok = jobs.cancel(task_id)
    if not ok:
        raise HTTPException(status_code=409, detail="not cancellable")
    return {"ok": True}