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, *, 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
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
_DEFAULT_TIMEOUT = 600.0
_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]) -> int:
    """Stream a response body to `dest` (atomic via .part). Returns byte count.
    Honors should_stop between chunks (partial file removed)."""
    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 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, *, timeout: float,
                should_stop: Callable[[], bool], fallback: str,
                headers: dict | None = None) -> FetchResult:
    resp = _http_get(url, timeout=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)
    return FetchResult(files=[dest], bytes=written)


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

def _fetch_dropbox(url: str, dest_dir: Path, *, 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, timeout=timeout,
                       should_stop=should_stop, fallback="dropbox-file")


def _fetch_pixeldrain(url: str, dest_dir: Path, *, 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, timeout=timeout,
                       should_stop=should_stop, fallback=f"{file_id}.bin")


def _fetch_mediafire(url: str, dest_dir: Path, *, timeout: float,
                     should_stop: Callable[[], bool]) -> FetchResult:
    page = _http_get(url, timeout=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, timeout=timeout,
                       should_stop=should_stop, fallback="mediafire-file")


def _fetch_gdrive(url: str, dest_dir: Path, *, timeout: float,
                  should_stop: Callable[[], bool]) -> FetchResult:
    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, *, timeout: float,
                should_stop: Callable[[], bool]) -> FetchResult:
    before = set(dest_dir.iterdir()) if dest_dir.exists() else set()
    _run_mega_get(url, str(dest_dir), timeout=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, *,
                   timeout: float = _DEFAULT_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, non-200, scrape miss, subprocess
    failure, stop) is captured on `.error` so the worker can record it and move
    on."""
    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, timeout=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}")