Blame for plugins/mega.py

import logging
import os
import re
import secrets
import time
from typing import Any, List, Optional, Tuple, Union
import requests
from Crypto.Cipher import AES
from Crypto.Util import Counter
from mega import Mega
from mega.crypto import (
    a32_to_str,
    base64_to_a32,
    base64_url_decode,
    decrypt_attr,
    decrypt_key,
)

from plugins.common import DownloadCancelled

logger = logging.getLogger("plugins.mega")

REQUEST_TIMEOUT = (15, 60)
MAX_FOLDER_FILES = 50


def is_mega_folder_url(url: str) -> bool:
    """Verifica si la URL corresponde a una carpeta de Mega.nz."""
    u = str(url or "").strip().lower()
    return "/folder/" in u or "/#f!" in u or "#f!" in u


def parse_mega_folder_url(url: str) -> Tuple[str, str]:
    """Extrae (folder_handle, folder_key) de URLs de carpetas de Mega."""
    u = str(url or "").strip()
    if "/folder/" in u:
        parts = u.split("/folder/", 1)[1].split("?", 1)[0].rstrip("/")
        if "#" in parts:
            handle, key = parts.split("#", 1)
            handle = handle.split("/")[0].strip()
            key = key.split("/")[0].strip()
            if handle and key:
                return handle, key
    for marker in ("/#F!", "/#f!", "#F!", "#f!"):
        if marker in u:
            parts = u.split(marker, 1)[1].split("?", 1)[0].rstrip("/")
            sub = parts.split("!")
            if len(sub) >= 2:
                handle = sub[0].strip()
                key = sub[1].strip()
                if handle and key:
                    return handle, key
    raise ValueError("El enlace no tiene un formato de carpeta de Mega reconocido.")


def parse_mega_file_url(url: str) -> Tuple[str, str]:
    """Extrae (file_handle, file_key) de URLs de archivos individuales de Mega."""
    u = str(url or "").strip()
    if "/file/" in u:
        parts = u.split("/file/", 1)[1].split("?", 1)[0].rstrip("/")
        if "#" in parts:
            handle, key = parts.split("#", 1)
            handle = handle.split("/")[0].strip()
            key = key.split("/")[0].strip()
            if handle and key:
                return handle, key
        elif "!" in parts:
            sub = parts.lstrip("#!").split("!")
            if len(sub) >= 2:
                handle = sub[0].strip()
                key = sub[1].strip()
                if handle and key:
                    return handle, key
    for marker in ("/#!", "#!"):
        if marker in u:
            parts = u.split(marker, 1)[1].split("?", 1)[0].rstrip("/")
            sub = parts.split("!")
            if len(sub) >= 2:
                handle = sub[0].strip()
                key = sub[1].strip()
                if handle and key:
                    return handle, key
    raise ValueError("El enlace no tiene un formato de archivo de Mega reconocido.")


def sanitize_filename(filename: str, default: str = "archivo_mega") -> str:
    """Limpia caracteres inválidos del nombre de archivo."""
    if not filename:
        return default
    clean = re.sub(r'[\/\\:\*\?"<>\|\x00-\x1f]', '_', filename).strip(' ._')
    return clean if clean else default


def _download_stream_decrypt(
    file_url: str,
    output_path: str,
    declared_size: int,
    k_str: bytes = None,
    initial_counter_val: int = None,
    aes_cipher: Optional[Any] = None,
    progress_fn=None,
    cancel_event=None,
    chunk_size: int = 512 * 1024,
    current_offset: int = 0,
    total_batch_size: int = 0,
    max_retries: int = 8,
) -> str:
    """
    Descarga por HTTP stream y descifra en tiempo real con AES-CTR directamente al archivo destino,
    con soporte de reanudación automática (Range: bytes=N-) y reintentos ante caídas de conexión.
    """
    target_dir = os.path.dirname(output_path)
    if target_dir:
        os.makedirs(target_dir, exist_ok=True)

    part_path = output_path + ".part"
    downloaded = 0
    if os.path.exists(part_path):
        cur_sz = os.path.getsize(part_path)
        # Reanudar estrictamente en múltiplos de bloques AES (16 bytes)
        downloaded = (cur_sz // 16) * 16
        if downloaded < cur_sz:
            try:
                with open(part_path, "r+b") as f:
                    f.truncate(downloaded)
            except Exception:
                downloaded = 0
                try:
                    os.remove(part_path)
                except Exception:
                    pass

    attempt = 0
    sess = requests.Session()

    while True:
        if cancel_event and cancel_event.is_set():
            raise DownloadCancelled()

        if declared_size > 0 and downloaded >= declared_size:
            break
        if declared_size == 0 and downloaded > 0:
            break

        attempt += 1
        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 downloaded > 0:
            headers["Range"] = f"bytes={downloaded}-"

        try:
            r = sess.get(file_url, headers=headers, stream=True, timeout=REQUEST_TIMEOUT)
            if r.status_code == 509:
                raise RuntimeError("Límite de transferencia de Mega excedido para la IP del worker.")
            if r.status_code not in (200, 206):
                r.raise_for_status()

            # Si el servidor no soportó Range y devolvió 200 en vez de 206 al reanudar, reiniciar desde 0
            if downloaded > 0 and r.status_code == 200:
                logger.warning("[Mega] El servidor no respetó Range en reanudación, reiniciando desde 0...")
                downloaded = 0
                if os.path.exists(part_path):
                    try:
                        os.remove(part_path)
                    except Exception:
                        pass

            # Inicializar cifrador AES-CTR en el offset correspondiente
            if k_str is not None and initial_counter_val is not None:
                block_offset = downloaded // 16
                cur_counter_val = initial_counter_val + block_offset
                counter = Counter.new(128, initial_value=cur_counter_val)
                active_aes = AES.new(k_str, AES.MODE_CTR, counter=counter)
            else:
                active_aes = aes_cipher

            mode = "ab" if (downloaded > 0 and os.path.exists(part_path)) else "wb"
            with open(part_path, mode, buffering=4 * 1024 * 1024) as f:
                for chunk in r.iter_content(chunk_size=chunk_size):
                    if cancel_event and cancel_event.is_set():
                        raise DownloadCancelled()
                    if chunk:
                        decrypted_chunk = active_aes.decrypt(chunk)
                        f.write(decrypted_chunk)
                        downloaded += len(chunk)
                        if progress_fn:
                            if total_batch_size > 0:
                                progress_fn(current_offset + downloaded, total_batch_size)
                            else:
                                progress_fn(downloaded, declared_size)

            if declared_size > 0 and downloaded >= declared_size:
                break
            if declared_size == 0 and downloaded > 0:
                break

            # Si el iter_content finalizó sin error pero no descargó todos los bytes
            if declared_size > 0 and downloaded < declared_size:
                logger.warning(f"[Mega] Stream finalizó prematuramente ({downloaded}/{declared_size} bytes). Reintentando...")
                if attempt >= max_retries:
                    raise RuntimeError(f"Descarga incompleta de Mega: se recibieron {downloaded} de {declared_size} bytes tras {max_retries} intentos.")
                time.sleep(min(2 * attempt, 10))
                continue

        except DownloadCancelled:
            raise
        except RuntimeError as re_err:
            if "Límite de transferencia" in str(re_err):
                raise
            if attempt >= max_retries:
                raise
            time.sleep(min(2 * attempt, 10))
        except Exception as conn_err:
            logger.warning(f"[Mega] Error en stream (intento {attempt}/{max_retries}): {conn_err}")
            if attempt >= max_retries:
                raise RuntimeError(f"Error descargando desde Mega tras {max_retries} intentos: {conn_err}")
            time.sleep(min(2 * attempt, 10))

    if declared_size > 0 and downloaded < declared_size:
        raise RuntimeError(f"Descarga incompleta de Mega: se recibieron {downloaded} de {declared_size} bytes.")
    if downloaded == 0:
        raise RuntimeError("No se descargaron datos de Mega (archivo vacío o servidor no respondió).")

    if not os.path.exists(part_path) or os.path.getsize(part_path) == 0:
        raise RuntimeError("El archivo descargado de Mega quedó vacío en disco.")

    if os.path.exists(output_path):
        try:
            os.remove(output_path)
        except Exception:
            pass
    os.replace(part_path, output_path)
    return output_path


def _download_mega_folder(
    url: str,
    output_dir: str,
    progress_fn=None,
    cancel_event=None
) -> List[str]:
    """Descarga e inspecciona recursivamente una carpeta pública de Mega."""
    if cancel_event and cancel_event.is_set():
        raise DownloadCancelled()

    folder_handle, folder_key = parse_mega_folder_url(url)

    rand_id = secrets.randbelow(2**31)
    resp = requests.post(
        "https://g.api.mega.co.nz/cs",
        params={"id": rand_id, "n": folder_handle},
        json=[{"a": "f", "c": 1, "r": 1}],
        timeout=REQUEST_TIMEOUT,
    )
    resp.raise_for_status()
    data = resp.json()
    if not isinstance(data, list) or not data or not isinstance(data[0], dict):
        raise RuntimeError("Mega rechazó la consulta de la carpeta.")

    nodes = data[0].get("f", [])
    if not nodes:
        raise RuntimeError("La carpeta de Mega está vacía o no contiene archivos descargables.")

    master_key_a32 = base64_to_a32(folder_key)
    file_items = []
    total_folder_bytes = 0

    for node in nodes:
        if node.get("t") != 0:  # 0 = archivo, 1 = carpeta
            continue
        k_val = node.get("k", "")
        if ":" not in k_val:
            continue
        k_enc = k_val.split(":")[1]
        try:
            k_a32 = base64_to_a32(k_enc)
            node_key_a32 = decrypt_key(k_a32, master_key_a32)
            k_derived = (
                node_key_a32[0] ^ node_key_a32[4],
                node_key_a32[1] ^ node_key_a32[5],
                node_key_a32[2] ^ node_key_a32[6],
                node_key_a32[3] ^ node_key_a32[7],
            )
            raw_attr = base64_url_decode(node.get("a", ""))
            attr = decrypt_attr(raw_attr, k_derived)
            if not attr or "n" not in attr:
                continue
            fn = sanitize_filename(attr["n"])
            size = int(node.get("s", 0))
            file_items.append({
                "handle": node.get("h"),
                "name": fn,
                "size": size,
                "key_a32": node_key_a32,
                "k_derived": k_derived,
            })
            total_folder_bytes += size
        except Exception as e:
            logger.debug(f"[Mega] No se pudo desencriptar nodo {node.get('h')}: {e}")

    if not file_items:
        raise RuntimeError("La carpeta de Mega está vacía o no contiene archivos descargables.")

    if len(file_items) > MAX_FOLDER_FILES:
        logger.warning(f"[Mega] Carpeta con {len(file_items)} archivos; limitando a {MAX_FOLDER_FILES}.")
        file_items = file_items[:MAX_FOLDER_FILES]

    downloaded_files = []
    overall_downloaded = 0

    for idx, item in enumerate(file_items, 1):
        if cancel_event and cancel_event.is_set():
            raise DownloadCancelled()

        f_handle = item["handle"]
        fn = item["name"]
        f_size = item["size"]
        node_key_a32 = item["key_a32"]
        k_derived = item["k_derived"]

        rand_id = secrets.randbelow(2**31)
        r_link = requests.post(
            "https://g.api.mega.co.nz/cs",
            params={"id": rand_id, "n": folder_handle},
            json=[{"a": "g", "g": 1, "n": f_handle}],
            timeout=REQUEST_TIMEOUT,
        )
        r_link.raise_for_status()
        link_data = r_link.json()
        if not isinstance(link_data, list) or not link_data or not isinstance(link_data[0], dict) or "g" not in link_data[0]:
            logger.warning(f"No se pudo obtener enlace de descarga para {fn} ({f_handle})")
            continue

        dl_url = link_data[0]["g"]
        out_file = os.path.join(output_dir, fn)

        k_str = a32_to_str(k_derived)
        iv = node_key_a32[4:6] + (0, 0)
        initial_counter_val = ((iv[0] << 32) + iv[1]) << 64

        _download_stream_decrypt(
            dl_url,
            out_file,
            f_size,
            k_str=k_str,
            initial_counter_val=initial_counter_val,
            progress_fn=progress_fn,
            cancel_event=cancel_event,
            current_offset=overall_downloaded,
            total_batch_size=total_folder_bytes,
        )

        overall_downloaded += f_size
        downloaded_files.append(out_file)

    if not downloaded_files:
        raise RuntimeError("No se pudo descargar ningún archivo de la carpeta de Mega.")

    return downloaded_files


def _download_mega_file(
    url: str,
    output_dir: str,
    progress_fn=None,
    cancel_event=None
) -> str:
    """Descarga un archivo individual de Mega con desencriptado en streaming al vuelo."""
    if cancel_event and cancel_event.is_set():
        raise DownloadCancelled()

    file_id, file_key = parse_mega_file_url(url)

    rand_id = secrets.randbelow(2**31)
    resp = requests.post(
        "https://g.api.mega.co.nz/cs",
        params={"id": rand_id},
        json=[{"a": "g", "g": 1, "p": file_id}],
        timeout=REQUEST_TIMEOUT,
    )
    resp.raise_for_status()
    data = resp.json()
    if not isinstance(data, list) or not data:
        raise RuntimeError("Mega rechazó la consulta del archivo.")

    res = data[0]
    if isinstance(res, int):
        if res == -9:
            raise RuntimeError("El archivo de Mega no existe o fue eliminado.")
        elif res == -16:
            raise RuntimeError("El archivo de Mega fue bloqueado por copyright o violación de términos.")
        elif res == -18:
            raise RuntimeError("Límite de transferencia de Mega excedido para la IP del worker.")
        elif res == -3:
            raise RuntimeError("Mega está ocupado temporalmente. Reintenta más tarde.")
        else:
            raise RuntimeError(f"Mega retornó código de error {res}.")

    if not isinstance(res, dict) or "g" not in res:
        raise RuntimeError("Mega no devolvió enlace de descarga directo.")

    file_url = res["g"]
    declared_size = int(res.get("s", 0))

    file_key_decoded = base64_to_a32(file_key)
    k = (
        file_key_decoded[0] ^ file_key_decoded[4],
        file_key_decoded[1] ^ file_key_decoded[5],
        file_key_decoded[2] ^ file_key_decoded[6],
        file_key_decoded[3] ^ file_key_decoded[7],
    )
    iv = file_key_decoded[4:6] + (0, 0)
    try:
        raw_attr = base64_url_decode(res.get("at", ""))
        attr = decrypt_attr(raw_attr, k)
        filename = attr["n"] if (isinstance(attr, dict) and "n" in attr) else f"mega_{file_id}"
    except Exception:
        filename = f"mega_{file_id}"

    safe_name = sanitize_filename(filename)
    output_path = os.path.join(output_dir, safe_name)

    k_str = a32_to_str(k)
    initial_counter_val = ((iv[0] << 32) + iv[1]) << 64

    return _download_stream_decrypt(
        file_url,
        output_path,
        declared_size,
        k_str=k_str,
        initial_counter_val=initial_counter_val,
        progress_fn=progress_fn,
        cancel_event=cancel_event,
    )


def download_mega(
    url: str,
    output_dir: str,
    progress_fn=None,
    cancel_event=None
) -> Union[str, List[str]]:
    """Descarga enlaces de MEGA (archivo individual o carpeta completa) con streaming directo y desencriptado al vuelo."""
    if cancel_event and cancel_event.is_set():
        raise DownloadCancelled()

    os.makedirs(output_dir, exist_ok=True)

    if is_mega_folder_url(url):
        logger.info(f"[Mega] Detectada carpeta de Mega: {url}")
        return _download_mega_folder(url, output_dir, progress_fn=progress_fn, cancel_event=cancel_event)

    # Archivo individual: intentar descarga nativa directa
    try:
        logger.info(f"[Mega] Detectado archivo individual de Mega: {url}")
        return _download_mega_file(url, output_dir, progress_fn=progress_fn, cancel_event=cancel_event)
    except DownloadCancelled:
        raise
    except Exception as native_err:
        logger.warning(f"[Mega] Falló descarga nativa directa ({native_err}), intentando fallback con mega.py...")
        try:
            mega = Mega()
            m = mega.login_anonymous()
            m = mega if m is None else m
            res = m.download_url(url, dest_path=output_dir)
            if not res or not os.path.exists(str(res)) or os.path.getsize(str(res)) == 0:
                if res and os.path.exists(str(res)):
                    try:
                        os.remove(str(res))
                    except Exception:
                        pass
                raise RuntimeError(str(native_err))
            return str(res)
        except DownloadCancelled:
            raise
        except Exception:
            raise native_err