"""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}")