configs/mega_folder_dl.py

#!/usr/bin/env python3
"""Paralelo: baixa um link publico de PASTA do mega.nz.

megadl e sequencial (~0,6 MB/s aqui). Varias conexoes contornam o throttle
por-conexao do Mega. Pula arquivos ja completos (compara tamanho).
"""
from __future__ import annotations
import argparse, base64, json, os, struct, sys, threading, time, urllib.request
from concurrent.futures import ThreadPoolExecutor
from Crypto.Cipher import AES
from Crypto.Util import Counter

API = "https://g.api.mega.co.nz/cs"
_seq = [0]
_lock = threading.Lock()


def b64d(s: str) -> bytes:
    s += "=" * (-len(s) % 4)
    return base64.b64decode(s.replace("-", "+").replace("_", "/"))


def api(payload, folder_id, tries=6):
    for attempt in range(tries):
        with _lock:
            _seq[0] += 1
            sid = _seq[0]
        req = urllib.request.Request(
            f"{API}?id={sid}&n={folder_id}",
            data=json.dumps(payload).encode(),
            headers={"Content-Type": "application/json"},
        )
        try:
            out = json.loads(urllib.request.urlopen(req, timeout=90).read())
        except Exception:
            time.sleep(2 * (attempt + 1)); continue
        # Um inteiro negativo no lugar do objeto e um erro do Mega (-3 = tente de novo).
        if isinstance(out, int):
            out = [out]
        if isinstance(out[0], int):
            if out[0] in (-3, -4):
                time.sleep(3 * (attempt + 1)); continue
            raise RuntimeError(f"Mega API error {out[0]}")
        return out
    raise RuntimeError("Mega API: falhou apos varias tentativas")


def node_keys(k_field: str, master: bytes):
    """O campo `k` pode trazer varias entradas `handle:chave` separadas por `/`
    (a pasta e compartilhada por mais de um no). So uma decifra com a master key
    do link, entao devolvemos todas as candidatas e quem chama escolhe a que
    produz um atributo valido."""
    ecb = AES.new(master, AES.MODE_ECB)
    out = []
    for entry in k_field.split("/"):
        if ":" not in entry:
            continue
        enc = b64d(entry.split(":", 1)[1])
        if len(enc) < 16:
            continue
        out.append(ecb.decrypt(enc[: len(enc) // 16 * 16]))
    return out


def node_key(k_field: str, master: bytes) -> bytes:
    keys = node_keys(k_field, master)
    if not keys:
        raise ValueError("no key entry")
    return keys[0]


def file_key(nk: bytes):
    w = struct.unpack(">8I", nk[:32])
    key = struct.pack(">4I", w[0] ^ w[4], w[1] ^ w[5], w[2] ^ w[6], w[3] ^ w[7])
    iv = struct.pack(">2I", w[4], w[5])  # 8 bytes de nonce
    return key, iv


def attr_name(attr_b64: str, key: bytes) -> str | None:
    raw = b64d(attr_b64)
    dec = AES.new(key, AES.MODE_CBC, b"\0" * 16).decrypt(raw[: len(raw) // 16 * 16])
    dec = dec.rstrip(b"\0")
    if not dec.startswith(b"MEGA"):
        return None
    try:
        return json.loads(dec[4:].decode("utf-8", "replace")).get("n")
    except Exception:
        return None


def download(node, dest_root, folder_id, master, state):
    handle, key, iv, size, relpath = node
    out = os.path.join(dest_root, relpath)
    if os.path.exists(out) and os.path.getsize(out) == size:
        with _lock:
            state["skipped"] += 1
        return
    os.makedirs(os.path.dirname(out), exist_ok=True)
    tmp = out + ".part"
    url = api([{"a": "g", "g": 1, "n": handle}], folder_id)[0]["g"]
    ctr = Counter.new(128, initial_value=int.from_bytes(iv + b"\0" * 8, "big"))
    aes = AES.new(key, AES.MODE_CTR, counter=ctr)
    with urllib.request.urlopen(url, timeout=180) as r, open(tmp, "wb") as f:
        while True:
            chunk = r.read(1 << 20)
            if not chunk:
                break
            f.write(aes.decrypt(chunk))
            with _lock:
                state["bytes"] += len(chunk)
    os.replace(tmp, out)
    with _lock:
        state["done"] += 1


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("link")
    ap.add_argument("--dest", required=True)
    ap.add_argument("--workers", type=int, default=12)
    a = ap.parse_args()

    folder_id, key_b64 = a.link.split("/folder/")[1].split("#")
    master = b64d(key_b64)
    nodes = api([{"a": "f", "c": 1, "r": 1, "ca": 1}], folder_id)[0]["f"]

    names, parents = {}, {}
    for n in nodes:
        parents[n["h"]] = n.get("p")
        if n["t"] == 1:
            names[n["h"]] = next(
                (nm for nk in node_keys(n["k"], master)
                 if (nm := attr_name(n["a"], nk[:16]))), n["h"])

    # Handle da pasta raiz do link: seu nome nao entra no caminho local
    # (megadl tambem o omite; manter igual permite retomar o que ele ja baixou).
    roots = {h for h in names if parents.get(h) not in names}

    def relpath_of(h):
        parts = []
        while h in names and h not in roots:
            parts.append(names[h]); h = parents.get(h)
        return os.path.join(*reversed(parts)) if parts else ""

    jobs, undecodable = [], []
    for n in nodes:
        if n["t"] != 0:
            continue
        chosen = None
        for nk in node_keys(n["k"], master):
            if len(nk) < 32:
                continue
            key, iv = file_key(nk)
            name = attr_name(n["a"], key)
            if name:
                chosen = (key, iv, name)
                break
        if chosen is None:
            undecodable.append(n["h"])
            continue
        key, iv, name = chosen
        jobs.append((n["h"], key, iv, n["s"], os.path.join(relpath_of(n["p"]), name)))
    if undecodable:
        print(f"AVISO: {len(undecodable)} nos sem chave decifravel, ignorados", flush=True)

    total = sum(j[3] for j in jobs)
    print(f"{len(jobs)} arquivos, {total/1e9:.2f} GB, {a.workers} conexoes", flush=True)
    state = {"bytes": 0, "done": 0, "skipped": 0}
    t0 = time.time()

    stop = threading.Event()
    def ticker():
        while not stop.wait(30):
            el = time.time() - t0
            print(f"[{el:6.0f}s] ok={state['done']} skip={state['skipped']}/{len(jobs)} "
                  f"{state['bytes']/1e6:8.0f} MB  {state['bytes']/1e6/el:5.2f} MB/s", flush=True)
    threading.Thread(target=ticker, daemon=True).start()

    errors = []
    with ThreadPoolExecutor(max_workers=a.workers) as ex:
        futs = {ex.submit(download, j, a.dest, folder_id, master, state): j for j in jobs}
        for f, j in futs.items():
            try:
                f.result()
            except Exception as e:
                errors.append((j[4], repr(e)))
    stop.set()
    print(f"FIM em {time.time()-t0:.0f}s  ok={state['done']} skip={state['skipped']} err={len(errors)}", flush=True)
    for p, e in errors[:40]:
        print("ERRO", p, e, flush=True)
    sys.exit(1 if errors else 0)


if __name__ == "__main__":
    main()