FabledCurator/backend/app/services/external_fetch.py

"""Fetchers for off-platform file hosts (mega / gdrive / mediafire / dropbox /
pixeldrain).

A shared, reusable subsystem: given an external_link URL, fetch the file(s) into
a destination directory and report the outcome. The download worker (separate
slice) drives these off the external_link ledger; any in-house downloader can
call `fetch_external()` directly.

No single tool covers all five hosts, so a small registry maps host → fetch
function behind one signature:

    fetch_external(host, url, dest_dir, *, read_timeout, total_timeout, should_stop)
        -> FetchResult

Backends:
  - dropbox    : force the direct-download variant (dl=1) + stream GET.
  - pixeldrain : GET the /api/file/{id} endpoint.
  - mediafire  : scrape the download page for the direct link + stream GET.
  - gdrive     : gdown (handles the confirm-token + virus-scan interstitial).
  - mega       : `megatools dl` subprocess (public link incl. #key); needs the
                 `megatools` binary in the runtime image (Debian apt package).

Public links work credential-free (rule 26); per-host creds are a later Settings
concern. Plain-HTTP homelab — no secure-context API. The HTTP / gdown /
subprocess calls go through module-level seams so unit tests run without
network, gdown, or MEGAcmd.
"""
from __future__ import annotations

import logging
import os
import re
import subprocess
import time
from collections.abc import Callable
from dataclasses import dataclass, field
from pathlib import Path
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit

import requests

log = logging.getLogger(__name__)

_CHUNK = 1 << 16
# Two distinct limits, because conflating them (the old single 3000s value) meant
# a stalled HTTP connection tied up a download-worker slot + the per-host lock for
# ~50 min before failing (operator-flagged 2026-06-17):
#   * READ timeout — max idle gap between bytes on an HTTP socket. A stalled host
#     (socket open, nothing flowing) is the common failure mode; a short read
#     timeout fails it fast. This is what requests' `timeout` actually enforces —
#     per-read, never a total.
#   * TOTAL budget — generous wall-clock cap for a file that IS actively
#     transferring (big films/packs). Enforced as a deadline across chunks, since
#     no HTTP client timeout bounds the total. Also the subprocess total for mega.
_CONNECT_TIMEOUT = 30.0
_READ_TIMEOUT = 60.0
_TOTAL_TIMEOUT = 1800.0  # 30 min per fetch
_USER_AGENT = (
    "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
    "(KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36"
)
# MediaFire's download page embeds the real file URL in a download-button href.
_MEDIAFIRE_RE = re.compile(
    r'href="(https://download[^"]+?\.mediafire\.com/[^"]+)"', re.IGNORECASE
)
_CD_FILENAME_RE = re.compile(r'filename\*?=(?:UTF-8\'\')?"?([^";]+)"?', re.IGNORECASE)


@dataclass
class FetchResult:
    files: list[Path] = field(default_factory=list)
    bytes: int = 0
    error: str | None = None

    @property
    def ok(self) -> bool:
        return self.error is None and bool(self.files)


class ExternalFetchError(Exception):
    """A fetch failed in a way worth recording on the link's last_error."""


# -- seams (monkeypatched in tests) ----------------------------------------

def _http_get(url: str, *, timeout: float, headers: dict | None = None,
              stream: bool = True) -> requests.Response:
    hdrs = {"User-Agent": _USER_AGENT}
    if headers:
        hdrs.update(headers)
    return requests.get(url, timeout=timeout, headers=hdrs, stream=stream)


def _gdown_download(url: str, out_dir: str) -> str | None:
    """Download a Google Drive url into out_dir via gdown; return the written
    path. Imported lazily so the dep is optional at module-import time."""
    # Lazy import: keeps gdown an optional dep that isn't needed to import this
    # module (e.g. the no-DB unit lane that never exercises a real fetch).
    import gdown
    # A trailing sep tells gdown to keep the server-side filename inside out_dir.
    return gdown.download(url, output=out_dir + os.sep, quiet=True, fuzzy=True)


def _run_mega_get(url: str, out_dir: str, *, timeout: float) -> None:
    """Download a mega.nz public link (key in the #fragment) into out_dir via
    `megatools dl` (the Debian `megatools` package). Raises ExternalFetchError
    on non-zero exit."""
    # Fixed argv (not shell): only `url` is external input, passed positionally,
    # so there's no shell-injection surface.
    proc = subprocess.run(
        ["megatools", "dl", "--path", out_dir, url],
        capture_output=True, text=True, timeout=timeout, check=False,
    )
    if proc.returncode != 0:
        raise ExternalFetchError(
            f"megatools dl exit {proc.returncode}: {(proc.stderr or '').strip()[:300]}"
        )


# -- helpers ---------------------------------------------------------------

def _safe_name(name: str, fallback: str) -> str:
    name = os.path.basename((name or "").strip().strip('"'))
    # Strip path separators / control chars; never empty.
    name = re.sub(r'[<>:"/\\|?*\x00-\x1f]', "_", name).strip(". ")
    return name or fallback


def _filename_from(resp: requests.Response, url: str, fallback: str) -> str:
    cd = resp.headers.get("Content-Disposition", "")
    m = _CD_FILENAME_RE.search(cd)
    if m:
        return _safe_name(m.group(1), fallback)
    path_name = os.path.basename(urlsplit(url).path)
    return _safe_name(path_name, fallback)


def _stream_to_file(resp: requests.Response, dest: Path,
                    should_stop: Callable[[], bool],
                    *, deadline: float | None = None) -> int:
    """Stream a response body to `dest` (atomic via .part). Returns byte count.
    Honors should_stop and the total-budget `deadline` (a time.monotonic() value)
    between chunks; the partial file is removed on either abort."""
    part = dest.with_name(dest.name + ".part")
    total = 0
    try:
        with part.open("wb") as fh:
            for chunk in resp.iter_content(chunk_size=_CHUNK):
                if should_stop():
                    raise ExternalFetchError("stopped")
                if deadline is not None and time.monotonic() > deadline:
                    raise ExternalFetchError("exceeded total fetch budget")
                if chunk:
                    fh.write(chunk)
                    total += len(chunk)
    except BaseException:
        part.unlink(missing_ok=True)
        raise
    os.replace(part, dest)
    return total


def _get_to_dir(url: str, dest_dir: Path, *, read_timeout: float,
                total_timeout: float, should_stop: Callable[[], bool],
                fallback: str, headers: dict | None = None) -> FetchResult:
    # (connect, read): a short read timeout fails a stalled socket fast; the total
    # budget is enforced separately as a deadline across chunks (requests has no
    # total-download timeout).
    resp = _http_get(
        url, timeout=(_CONNECT_TIMEOUT, read_timeout), headers=headers, stream=True
    )
    if resp.status_code != 200:
        return FetchResult(error=f"HTTP {resp.status_code} for {url}")
    name = _filename_from(resp, url, fallback)
    dest = dest_dir / name
    written = _stream_to_file(
        resp, dest, should_stop, deadline=time.monotonic() + total_timeout
    )
    return FetchResult(files=[dest], bytes=written)


# -- per-host fetchers -----------------------------------------------------

def _fetch_dropbox(url: str, dest_dir: Path, *, read_timeout: float,
                   total_timeout: float,
                   should_stop: Callable[[], bool]) -> FetchResult:
    # Force the direct-download variant: dl=1 (Dropbox serves an HTML preview
    # for dl=0). Rewrite/insert the param rather than string-replace so ?dl=0,
    # &dl=0, and a missing param all resolve.
    parts = urlsplit(url)
    q = dict(parse_qsl(parts.query))
    q["dl"] = "1"
    direct = urlunsplit(parts._replace(query=urlencode(q)))
    return _get_to_dir(direct, dest_dir, read_timeout=read_timeout,
                       total_timeout=total_timeout, should_stop=should_stop,
                       fallback="dropbox-file")


def _fetch_pixeldrain(url: str, dest_dir: Path, *, read_timeout: float,
                      total_timeout: float,
                      should_stop: Callable[[], bool]) -> FetchResult:
    # /u/{id} (and /l/{id}) → the API file endpoint.
    file_id = urlsplit(url).path.rstrip("/").split("/")[-1]
    if not file_id:
        return FetchResult(error=f"no pixeldrain id in {url}")
    api = f"https://pixeldrain.com/api/file/{file_id}"
    return _get_to_dir(api, dest_dir, read_timeout=read_timeout,
                       total_timeout=total_timeout, should_stop=should_stop,
                       fallback=f"{file_id}.bin")


def _fetch_mediafire(url: str, dest_dir: Path, *, read_timeout: float,
                     total_timeout: float,
                     should_stop: Callable[[], bool]) -> FetchResult:
    page = _http_get(url, timeout=(_CONNECT_TIMEOUT, read_timeout), stream=False)
    if page.status_code != 200:
        return FetchResult(error=f"HTTP {page.status_code} for mediafire page")
    m = _MEDIAFIRE_RE.search(page.text or "")
    if not m:
        return FetchResult(error="mediafire direct link not found on page")
    return _get_to_dir(m.group(1), dest_dir, read_timeout=read_timeout,
                       total_timeout=total_timeout, should_stop=should_stop,
                       fallback="mediafire-file")


def _fetch_gdrive(url: str, dest_dir: Path, *, read_timeout: float,
                  total_timeout: float,
                  should_stop: Callable[[], bool]) -> FetchResult:
    # gdown manages its own HTTP session/timeouts; the task's celery hard limit is
    # the outer backstop. read_timeout/total_timeout are accepted for a uniform
    # registry signature but not separately enforceable here.
    out = _gdown_download(url, str(dest_dir))
    if not out:
        return FetchResult(error="gdown returned no file (quota / private?)")
    p = Path(out)
    if not p.exists():
        return FetchResult(error=f"gdown reported {out} but it is missing")
    return FetchResult(files=[p], bytes=p.stat().st_size)


def _fetch_mega(url: str, dest_dir: Path, *, read_timeout: float,
                total_timeout: float,
                should_stop: Callable[[], bool]) -> FetchResult:
    before = set(dest_dir.iterdir()) if dest_dir.exists() else set()
    # megatools is a subprocess: its timeout IS a total wall-clock cap (the read
    # timeout has no analogue here), so the total budget applies directly.
    _run_mega_get(url, str(dest_dir), timeout=total_timeout)
    new = [p for p in dest_dir.iterdir() if p not in before and p.is_file()]
    if not new:
        return FetchResult(error="mega-get wrote no new file")
    return FetchResult(files=new, bytes=sum(p.stat().st_size for p in new))


_REGISTRY: dict[str, Callable[..., FetchResult]] = {
    "dropbox": _fetch_dropbox,
    "pixeldrain": _fetch_pixeldrain,
    "mediafire": _fetch_mediafire,
    "gdrive": _fetch_gdrive,
    "mega": _fetch_mega,
}

SUPPORTED_HOSTS = tuple(_REGISTRY)


def fetch_external(host: str, url: str, dest_dir: Path, *,
                   read_timeout: float = _READ_TIMEOUT,
                   total_timeout: float = _TOTAL_TIMEOUT,
                   should_stop: Callable[[], bool] = lambda: False) -> FetchResult:
    """Fetch `url` (a `host` link) into `dest_dir`. Returns a FetchResult; never
    raises — any backend error (transport, read/total timeout, non-200, scrape
    miss, subprocess failure, stop) is captured on `.error` so the worker can
    record it and move on.

    `read_timeout` fails a stalled HTTP socket fast (idle gap between bytes);
    `total_timeout` is the generous wall-clock cap for a large file that is
    actively transferring (and the subprocess total for mega)."""
    fetcher = _REGISTRY.get(host)
    if fetcher is None:
        return FetchResult(error=f"unsupported host {host!r}")
    dest_dir.mkdir(parents=True, exist_ok=True)
    try:
        return fetcher(url, dest_dir, read_timeout=read_timeout,
                       total_timeout=total_timeout, should_stop=should_stop)
    except requests.RequestException as exc:
        return FetchResult(error=f"transport error: {exc}")
    except subprocess.TimeoutExpired:
        return FetchResult(error="timed out")
    except ExternalFetchError as exc:
        return FetchResult(error=str(exc))
    except Exception as exc:  # never let a backend quirk kill the worker
        log.warning("external fetch (%s) failed for %s: %s", host, url, exc)
        return FetchResult(error=f"{type(exc).__name__}: {exc}")