THOR / plugins /mega.py
ElJoker63's picture
Actualizar worker
e4ae20c verified
Raw History Blame Contribute Delete
17.6 kB
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