ani-web/server/upstream.py

"""upstream.py — provider-agnostic source layer.

A *provider* is a module in providers/ implementing the same interface:

    search(q)                  -> [{id, slug_id, name, poster}]
    details(slug_id)           -> {id, name, synopsis, mal_id, poster, seasons}
    episodes(anime_id)         -> [{id, number, filler}]
    languages(ep_id)           -> [{code, name, embed_url}]
    master_m3u8(ep_id, lang)   -> master m3u8 url | None

Switching sources = add providers/<name>.py, then:
    ANI_SOURCE=other,anidb      # ordered fallback: first to succeed wins
(default ANI_SOURCE=hianime,anidb,anikoto). Exceptions trigger the fallback; a legitimate
None from master_m3u8 ("no dub") is final, not an error.

Shared HTTP helpers (_get/get_stream/fetch_parallel, Chrome TLS
impersonation, TTL cache) live here so every provider gets them for free;
they work on any URL, not just anidb's.
"""
import base64
import importlib
import json
import os
import re
import threading
import time

from Crypto.Cipher import AES
from curl_cffi import requests

H = dict(impersonate="chrome", timeout=20)

# ponytail: thread-local sessions (uvicorn threadpool); curl_cffi Sessions
# are not thread-safe, but per-thread keep-alive still reuses TCP/TLS across
# the many requests a worker handles. Saves a handshake on every call.
_tls = threading.local()

def _session():
    if not hasattr(_tls, "s"):
        _tls.s = requests.Session()
    return _tls.s

# ponytail: in-memory TTL cache shared by all providers (key namespace is
# the provider's job). No lock: worst case a duplicate fetch. A dead/rotated
# playlist URL is recovered by the one-shot clear_cache()+retry in api_stream.
_CACHE = {}  # key -> (expires, value)

def _cached(key, ttl, fn, *a, **kw):
    v = _CACHE.get(key)
    if v and v[0] > time.time():
        return v[1]
    val = fn(*a, **kw)
    _CACHE[key] = (time.time() + ttl, val)
    if len(_CACHE) > 256:
        _CACHE.clear()
    return val

def clear_cache():
    _CACHE.clear()

# MegaPlay's getSourcesNew hid the stream URL in an AES-256-CBC "enc" field.
# Key/IV are static and hardcoded in their public player JS (newclient.min.js)
# — anti-scrape theater, not access control. ponytail: breaks if they rotate
# the key; upgrade path = re-read newclient.min.js (or let it fail: both
# shapes are tolerated below, so the old code path still works if reverted).
_MP_KEY = b"i?LMTAx0Q6,:}50U".ljust(32, b"\x00")
_MP_IV = b"W0;27ToaUpl_P%'c"

def megaplay_sources(d):
    """getSourcesNew response -> same dict with sources["file"] filled.
    Old players return it plaintext; current ones base64url+AES-wrap it in
    "enc" (decrypted payload is {"file": ...})."""
    if not (d.get("sources") or {}).get("file") and d.get("enc"):
        raw = base64.urlsafe_b64decode(d["enc"] + "=" * (-len(d["enc"]) % 4))
        pt = AES.new(_MP_KEY, AES.MODE_CBC, _MP_IV).decrypt(raw)
        j = json.loads(pt[:-pt[-1]])  # strip PKCS7
        if j.get("file"):
            d = {**d, "sources": {"file": j["file"]}}
    return d


def megaplay_master(embed, referer):
    """MegaPlay embed url -> (master_url, [{label, src} subs], referer_or_None).

    Shared by every provider fronting the MegaPlay player (hianime,
    animekai): embed page (referer-gated) -> data-id -> getSourcesNew
    (megaplay_sources) -> prove the CDN serves us, remembering which
    Referer variant worked. referer = the front site's home (the embed
    page 403s a stub without it). One attempt: the enc payload is
    deterministic per data-id, re-rolling would just spam the provider.
    The caller registers the returned master's hosts in its own
    /stream allowlist."""
    page = _get(embed, headers={"Referer": referer}, timeout=15).text
    mid = re.search(r'data-id="(\d+)"', page)
    if not mid:
        return None, [], None
    host = re.match(r'https?://([^/]+)', embed).group(1)
    d = megaplay_sources(_get(f"https://{host}/stream/getSourcesNew?id={mid.group(1)}",
             headers={"Referer": embed, "X-Requested-With": "XMLHttpRequest"},
             timeout=15).json())
    src = (d.get("sources") or {}).get("file")
    if not src:
        return None, [], None
    subs = [{"label": (x or {}).get("label", ""), "src": (x or {}).get("file")}
            for x in d.get("tracks") or []]
    for cand, refr in (({"Referer": embed}, embed), ({}, None)):
        try:
            _get(src, headers=cand, timeout=10)
            return src, subs, refr
        except Exception:
            continue
    return None, [], None


def _get(url, **kw):
    # ponytail: one retry (300ms) on transient upstream errors (Cloudflare
    # 403/503, 429, 5xx); 404 and other 4xx are final. Failing fast matters:
    # hls.js's fragLoadingTimeOut (20s) races curl's timeout (20s), so a
    # quick 502 lets hls.js's own frag retries (6x, backoff) do the waiting
    for attempt in range(2):
        r = _session().get(url, **{**H, **kw})  # kw overrides H (e.g. timeout)
        if not (r.status_code in (403, 429, 503) or r.status_code >= 500
                or "Just a moment" in r.text[:500]):
            r.raise_for_status()
            return r
        if attempt:
            raise RuntimeError(f"upstream {r.status_code}: {url}")
        time.sleep(0.3)


def get_stream(url, timeout=12, headers=None):
    """Segments: open the upstream response WITHOUT reading the body, so the
    proxy can forward bytes as the CDN emits them. (A buffered fetch made the
    browser's time-to-first-byte = the CDN's whole upload time, which tripped
    hls.js's fragLoadTimeOut on slow CDN moments.)
    ponytail: skips the 'Just a moment' body check — challenge pages arrive
    as 403/503 (handled below), and reading the body would defeat streaming."""
    for attempt in range(2):
        r = _session().get(url, **{**H, "timeout": timeout, "stream": True,
                                  "headers": headers or {}})
        if r.status_code in (403, 429, 503) or r.status_code >= 500:
            r.close()
            if attempt:
                raise RuntimeError(f"upstream {r.status_code}: {url}")
            time.sleep(0.3)
            continue
        r.raise_for_status()
        return r


def fetch_parallel(url, lanes=16, lane=32 * 1024, timeout=12, headers=None):
    """One upstream file across N parallel byte-range lanes (16, like yt-dlp's
    -N 16 / ani-cli's download path).

    Per-connection CDN throttling is the common slow-segment cause; N lanes
    turn one 60s crawl into N concurrent crawls (I/O-bound, GIL irrelevant).
    Returns (body, content-type); raises on ANY failure so the caller falls
    back to the single-request path — never worse than before."""
    from concurrent.futures import ThreadPoolExecutor
    head = _session().head(url, **{**H, "timeout": timeout,
                                  "headers": headers or {}})
    head.raise_for_status()
    total = int(head.headers.get("content-length") or 0)
    ct = head.headers.get("content-type", "application/octet-stream")
    if total == 0:
        raise RuntimeError("no content-length")
    if total <= 2 * lane:          # too small to be worth N connections
        r = _get(url, timeout=timeout)
        return r.content, ct
    n = min(lanes, total // lane)
    size = (total + n - 1) // n
    parts = [None] * n

    def grab(i):
        a = i * size
        b = total - 1 if i == n - 1 else (i + 1) * size - 1
        r = _get(url, headers={**(headers or {}), "Range": f"bytes={a}-{b}"},
                 timeout=timeout)
        if r.status_code not in (200, 206) or len(r.content) != b - a + 1:
            raise RuntimeError(f"lane {i}: http {r.status_code} "
                               f"len={len(r.content)} want={b - a + 1}")
        parts[i] = r.content

    with ThreadPoolExecutor(max_workers=n) as ex:
        list(ex.map(grab, range(n)))
    body = b"".join(parts)
    if len(body) != total:
        raise RuntimeError(f"stitch mismatch {len(body)} != total")
    return body, ct


# ---------- Jikan intros: provider-independent (keyless MAL mirror) ----------

def _parse_intros(eps):
    """Jikan episode list -> one entry per episode: {title?, start?, length?}.
    title from MAL (episode name); start/length only when MAL has the timing."""
    out = []
    for e in eps:
        i = e.get("intro") or {}
        d = {"title": e.get("title") or None}
        if i.get("start") is not None:
            d.update(start=int(i["start"]), length=int(i["length"]))
        out.append(d)
    return out


def _intros(mal_id):
    """Jikan (keyless MAL mirror): per-episode intro start/length, seconds."""
    # Jikan v4 returns the list directly under 'data'; tolerate a wrapped
    # shape too, so a format change doesn't silently kill the button
    r = _get("https://api.jikan.moe/v4/anime/{}".format(mal_id) + "/episodes")
    d = r.json().get("data")
    eps = d if isinstance(d, list) else (d or {}).get("episodes", [])
    return _parse_intros(eps)


intros = lambda m: _cached(f"jikan:intros:{m}", 30 * 86400, _intros, m)  # static data


# ---------- source dispatch ----------

def _providers():
    # ponytail: env read per call, so tests and runtime can re-point without
    # re-importing; comma list = ordered fallback
    return [p.strip() for p in os.environ.get("ANI_SOURCE", "hianime,anidb,anikoto").split(",")
            if p.strip()]


def _prov(name):
    try:
        return importlib.import_module(f".providers.{name}", __package__)
    except Exception:
        return None


def _call(name, *a, **kw):
    errs = []
    for p in _providers():
        try:
            return getattr(_prov(p), name)(*a, **kw)
        except Exception as e:
            errs.append(f"{p}: {e}")
    raise RuntimeError("all sources failed -> " + " | ".join(errs))


search = lambda q: _call("search", q)
details = lambda s: _call("details", s)
episodes = lambda i: _call("episodes", i)
languages = lambda e: _call("languages", e)
master_m3u8 = lambda e, l="eng": _call("master_m3u8", e, l)


def stream_meta(url):
    """Provider hook: extra upstream headers for a /stream URL (hianime's
    CDN may demand the embed origin as Referer), or None if no enabled
    provider claims the host — that doubles as the allowlist (an empty
    dict means: allowed, no extra headers)."""
    for p in _providers():
        fn = getattr(_prov(p), "stream_meta", None)
        if fn:
            v = fn(url)
            if v is not None:
                return v
    return None


def subtitles(ep_id):
    """External .vtt url for an episode, or None (anidb bakes subs in)."""
    for p in _providers():
        fn = getattr(_prov(p), "subtitles", None)
        if fn:
            v = fn(ep_id)
            if v:
                return v
    return None


if __name__ == "__main__":
    # ponytail: assert self-check instead of a test framework.
    # Needs a live source: ANI_SOURCE=<provider> python -m server.upstream
    results = search("chainsaw man")
    assert results and any(x["name"] == "Chainsaw Man" for x in results), results
    top = next(x for x in results if x["name"] == "Chainsaw Man")
    d = details(top["slug_id"])
    assert d["mal_id"] == 44511 and d["synopsis"], d
    assert [s["id"] for s in d["seasons"]] == [922, 923], d["seasons"]
    eps = episodes(top["id"])
    assert eps and eps[0]["number"] == 1, eps
    m = master_m3u8(eps[0]["id"], "eng")
    assert m and "EXT-X-STREAM" in _get(m).text, m
    print(f"OK: {top['name']} (MAL {d['mal_id']}, {len(d['seasons'])} seasons, "
          f"{len(eps)} eps) -> ep1 dub: {m[:70]}...")
    import sys
    sys.exit(0)