Blame for plugins/stream.py

import logging
import os
import re
import socket
import time
from http.client import IncompleteRead as HttpIncompleteRead
from typing import Optional, Tuple
from urllib.parse import parse_qs, unquote, urlparse
import requests
from requests.exceptions import (
    ChunkedEncodingError,
    ConnectionError as ReqConnectionError,
    HTTPError,
    RequestException,
    Timeout as ReqTimeout,
)
from urllib3.exceptions import IncompleteRead as Urllib3IncompleteRead, ProtocolError

from plugins.common import DownloadCancelled

logger = logging.getLogger("hf_worker.plugins.stream")


def _is_master_stream_url(url: str) -> bool:
    """

    Indica si la URL apunta al streaming interno del Master (archivo puente o

    stream de Telegram), y no a un recurso externo.



    Se comprueba el host completo y la ruta, no un substring: una URL externa

    arbitraria que contenga "/api/cluster/" en el path o en el query (por

    ejemplo "https://otro.host/x?u=/api/cluster/y") NO debe clasificarse como

    interna, porque eso saltaría las validaciones anti-HTML de este módulo.

    """
    if not url:
        return False
    if url.startswith("tg://"):
        return True
    try:
        parsed = urlparse(url)
    except Exception:
        return False
    host = (parsed.netloc or "").lower().split("@")[-1].split(":")[0]
    master_host = (urlparse(os.getenv("MANAGER_URL", "") or "").netloc or "").lower().split(":")[0]
    if not master_host:
        master_host = "downdus.udyat.site"
    return host == master_host and (parsed.path or "").startswith("/api/cluster/")


def _extract_filename_from_headers_or_url(url: str, headers: dict, fallback_name: Optional[str] = None) -> str:
    """Extrae el nombre de archivo de Content-Disposition o de la URL."""
    cd = headers.get("Content-Disposition", "") or headers.get("content-disposition", "")
    if not cd and "?" in url:
        try:
            qs = parse_qs(urlparse(url).query)
            for k in ("response-content-disposition", "content-disposition"):
                if k in qs and qs[k]:
                    cd = qs[k][0]
                    break
        except Exception:
            pass

    if cd:
        match_utf8 = re.search(r"filename\*=UTF-8''([^;\r\n]+)", cd, re.IGNORECASE)
        if match_utf8:
            raw_fn = unquote(match_utf8.group(1)).strip().strip('"\'')
            clean = os.path.basename(raw_fn)
            clean = re.sub(r'[\/\\:\*\?"<>\|\x00-\x1f]', '_', clean).strip(' ._')
            if clean:
                name, ext = os.path.splitext(clean)
                return (name[:100] + ext) if len(name) > 100 else clean
        match_std = re.search(r'filename=["\']?([^"\';\r\n]+)["\']?', cd, re.IGNORECASE)
        if match_std:
            raw_fn = unquote(match_std.group(1)).strip()
            clean = os.path.basename(raw_fn)
            clean = re.sub(r'[\/\\:\*\?"<>\|\x00-\x1f]', '_', clean).strip(' ._')
            if clean:
                name, ext = os.path.splitext(clean)
                return (name[:100] + ext) if len(name) > 100 else clean

    try:
        path = urlparse(url).path
        base = os.path.basename(path)
        if base and "." in base:
            clean = unquote(base).strip()
            clean = os.path.basename(clean)
            clean = re.sub(r'[\/\\:\*\?"<>\|\x00-\x1f]', '_', clean).strip(' ._')
            if clean:
                name, ext = os.path.splitext(clean)
                return (name[:100] + ext) if len(name) > 100 else clean
    except Exception:
        pass

    if fallback_name:
        clean_fallback = os.path.basename(str(fallback_name)).strip()
        clean_fallback = re.sub(r'[\/\\:\*\?"<>\|\x00-\x1f]', '_', clean_fallback).strip(' ._')
        if clean_fallback and clean_fallback not in ("downloaded_file", "archivo"):
            return clean_fallback

    return "archivo_descargado"


def _safe_close(resp):
    """Cierra la respuesta HTTP de forma segura."""
    if resp is not None and hasattr(resp, "close") and callable(resp.close):
        try:
            resp.close()
        except Exception:
            pass


def _make_request(sess: requests.Session, req_url: str, req_headers: dict) -> requests.Response:
    """Ejecuta la petición HTTP con pool y fallback automático a cloudscraper si hay bloqueos."""
    r = None
    try:
        r = sess.get(req_url, headers=req_headers, stream=True, timeout=(30, 90), allow_redirects=True)
        if r.status_code in (403, 503):
            try:
                import cloudscraper
                scraper = cloudscraper.create_scraper()
                sc_r = scraper.get(req_url, headers=req_headers, stream=True, timeout=(30, 90), allow_redirects=True)
                if sc_r.status_code < 400:
                    _safe_close(r)
                    r = sc_r
            except Exception:
                pass
        return r
    except Exception as req_err:
        try:
            import cloudscraper
            scraper = cloudscraper.create_scraper()
            return scraper.get(req_url, headers=req_headers, stream=True, timeout=(30, 90), allow_redirects=True)
        except Exception:
            raise req_err


_stream_session = None


def get_stream_session() -> requests.Session:
    """Sesión HTTP persistente con pool de conexiones y keep-alive para alta velocidad."""
    global _stream_session
    if _stream_session is None:
        _stream_session = requests.Session()
        adapter = requests.adapters.HTTPAdapter(
            pool_connections=20,
            pool_maxsize=40,
            max_retries=requests.adapters.Retry(
                total=3,
                backoff_factor=0.5,
                status_forcelist=[500, 502, 503, 504],
                raise_on_status=False,
            ),
        )
        _stream_session.mount("http://", adapter)
        _stream_session.mount("https://", adapter)
    return _stream_session


def download_stream_url(

    url: str,

    output_path: str,

    progress_fn=None,

    cancel_event=None,

    headers: dict = None,

    chunk_size: int = 256 * 1024,

    max_retries: int = 10,

    custom_filename: Optional[str] = None,

) -> Tuple[str, int]:
    """Descarga por HTTP/HTTPS en streaming reportando progreso en chunks con socket y pool optimizados,

    incorporando reanudación automática (Range: bytes=N-) ante desconexiones y microcortes de red."""
    if cancel_event and cancel_event.is_set():
        raise DownloadCancelled()

    req_headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36",
        "Accept": "*/*",
    }
    if "apkmirror.com" in url.lower():
        req_headers.update({
            "Referer": "https://www.apkmirror.com/",
            "Origin": "https://www.apkmirror.com",
            "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,*/*;q=0.8",
        })
    elif "apkvision." in url.lower():
        req_headers.update({
            "Referer": "https://apkvision.org/",
            "Origin": "https://apkvision.org",
        })
    if headers:
        req_headers.update(headers)

    sess = get_stream_session()
    started = time.perf_counter()

    r = _make_request(sess, url, req_headers)
    r.raise_for_status()

    total_size = int(r.headers.get("Content-Length", 0) or r.headers.get("content-length", 0))
    ctype = (r.headers.get("Content-Type") or r.headers.get("content-type") or "").lower()

    real_fn = _extract_filename_from_headers_or_url(r.url, r.headers, fallback_name=custom_filename)

    if "://" in output_path or output_path.startswith(("http://", "https://")):
        parts = re.split(r'https?://', output_path)
        clean_dir = parts[0].rstrip("/\\") if parts and parts[0] else ""
        if clean_dir and os.path.isdir(clean_dir):
            target_dir = clean_dir
        else:
            target_dir = os.path.dirname(output_path) if os.path.isdir(os.path.dirname(output_path)) else "/tmp"
        output_path = os.path.join(target_dir, real_fn)
    elif os.path.isdir(output_path) or output_path.endswith(os.sep) or output_path.endswith("/"):
        target_dir = output_path
        output_path = os.path.join(target_dir, real_fn)
    elif os.path.basename(output_path) in ("downloaded_file", "archivo"):
        target_dir = os.path.dirname(output_path)
        output_path = os.path.join(target_dir, real_fn)

    target_dir = os.path.dirname(output_path)
    if target_dir:
        os.makedirs(target_dir, exist_ok=True)

    is_internal_stream = _is_master_stream_url(url)
    allowed_doc_exts = {".html", ".htm", ".xhtml", ".xml", ".svg", ".txt", ".json", ".md", ".css", ".js"}

    # Validación anti-páginas de bloqueo/error HTML cuando se espera archivo multimedia
    expected_ext = os.path.splitext(output_path)[1].lower()
    if not is_internal_stream and expected_ext not in allowed_doc_exts and "text/html" in ctype and expected_ext in [".mp4", ".mp3", ".mkv", ".avi", ".zip", ".7z", ".rar", ".apk", ".bin"]:
        _safe_close(r)
        raise RuntimeError("El servidor no devolvió un archivo multimedia descargable.")

    part_path = output_path + ".part"
    if os.path.exists(part_path):
        try:
            os.remove(part_path)
        except Exception:
            pass

    current_resp = r
    current_url = getattr(r, "url", url) or url
    downloaded = 0
    attempt = 0

    while True:
        if cancel_event and cancel_event.is_set():
            _safe_close(current_resp)
            if os.path.exists(part_path):
                try:
                    os.remove(part_path)
                except Exception:
                    pass
            raise DownloadCancelled()

        try:
            # Si venimos de un corte de conexión, reanudamos con HTTP Range
            if current_resp is None:
                h = dict(req_headers)
                if downloaded > 0:
                    h["Range"] = f"bytes={downloaded}-"
                logger.info(
                    "Reanudando descarga: %s | Offset: %d bytes | Intento %d/%d",
                    current_url,
                    downloaded,
                    attempt,
                    max_retries,
                )
                try:
                    current_resp = _make_request(sess, current_url, h)
                except Exception:
                    if current_url != url:
                        logger.warning("Fallo conectando a URL derivada, reintentando con URL original: %s", url)
                        current_url = url
                        current_resp = _make_request(sess, current_url, h)
                    else:
                        raise

                # Manejo de respuesta 416 (Range Not Satisfiable)
                if current_resp.status_code == 416:
                    if total_size > 0 and downloaded >= total_size:
                        _safe_close(current_resp)
                        break
                    logger.warning("Error HTTP 416 recibido. Reiniciando descarga desde byte 0...")
                    downloaded = 0
                    h.pop("Range", None)
                    _safe_close(current_resp)
                    current_resp = _make_request(sess, current_url, h)

                current_resp.raise_for_status()

                if current_resp.status_code == 206:
                    file_mode = "ab"
                    cr = current_resp.headers.get("Content-Range", "") or current_resp.headers.get("content-range", "")
                    cr_m = re.search(r"/(\d+)", cr)
                    if cr_m and total_size <= 0:
                        total_size = int(cr_m.group(1))
                elif current_resp.status_code == 200:
                    if downloaded > 0:
                        logger.warning("El servidor no soporta Range (HTTP 200). Reiniciando descarga desde el inicio...")
                        downloaded = 0
                    file_mode = "wb"
                    cl = int(current_resp.headers.get("Content-Length", 0) or current_resp.headers.get("content-length", 0))
                    if cl > 0 and total_size <= 0:
                        total_size = cl
                else:
                    file_mode = "ab" if downloaded > 0 else "wb"
            else:
                file_mode = "wb"

            with open(part_path, file_mode, buffering=4 * 1024 * 1024) as f:
                for chunk in current_resp.iter_content(chunk_size=chunk_size):
                    if cancel_event and cancel_event.is_set():
                        raise DownloadCancelled()
                    if chunk:
                        f.write(chunk)
                        downloaded += len(chunk)
                        if progress_fn:
                            progress_fn(downloaded, total_size)

            _safe_close(current_resp)
            current_resp = None

            # Si el stream finalizó antes de tiempo pero faltan bytes, provocar reintento
            if total_size > 0 and downloaded < total_size:
                raise Urllib3IncompleteRead(downloaded, total_size - downloaded)

            # Descarga completada exitosamente
            break

        except DownloadCancelled:
            _safe_close(current_resp)
            if os.path.exists(part_path):
                try:
                    os.remove(part_path)
                except Exception:
                    pass
            raise

        except Exception as e:
            _safe_close(current_resp)
            current_resp = None

            if isinstance(e, HTTPError) and getattr(e, "response", None) is not None:
                status = e.response.status_code
                if status in (400, 401, 403, 404, 405, 410):
                    if os.path.exists(part_path):
                        try:
                            os.remove(part_path)
                        except Exception:
                            pass
                    raise

            attempt += 1
            if attempt > max_retries:
                logger.error("Descarga falló tras %d intentos. Último error: %s", max_retries, e)
                if os.path.exists(part_path):
                    try:
                        os.remove(part_path)
                    except Exception:
                        pass
                raise

            if os.path.exists(part_path):
                downloaded = os.path.getsize(part_path)

            sleep_time = min(2 * attempt, 10)
            logger.warning(
                "Conexión interrumpida durante descarga (%s: %s). Reintentando (%d/%d) en %ds desde byte %d/%d...",
                type(e).__name__,
                e,
                attempt,
                max_retries,
                sleep_time,
                downloaded,
                total_size,
            )
            time.sleep(sleep_time)

    # Renombrar atómicamente el archivo temporal .part al destino final
    if os.path.exists(output_path):
        try:
            os.remove(output_path)
        except Exception:
            pass
    os.replace(part_path, output_path)

    # Validación post-descarga: solo para descargas externas que no sean archivos de texto/web/documentos
    if not is_internal_stream and expected_ext not in allowed_doc_exts:
        if os.path.exists(output_path) and os.path.getsize(output_path) < 150 * 1024:
            is_html = False
            try:
                with open(output_path, "rb") as test_f:
                    head = test_f.read(512).lower()
                    if b"<!doctype html" in head or b"<html" in head:
                        is_html = True
            except Exception:
                pass

            if is_html:
                try:
                    os.remove(output_path)
                except Exception:
                    pass
                raise RuntimeError("El contenido descargado no es un archivo multimedia válido.")

    elapsed = max(time.perf_counter() - started, 0.001)
    if "/api/cluster/tg_stream/" in url:
        logger.info(
            "[TG-Stream Worker] recibido %s bytes a %.2f MiB/s en %.2fs (%s)",
            downloaded,
            downloaded / elapsed / (1024 * 1024),
            elapsed,
            os.path.basename(output_path),
        )
    return output_path, downloaded