540 lines
22 KiB
Python
540 lines
22 KiB
Python
#!/usr/bin/env python3
|
|
"""Warhammer DLP Helper.
|
|
|
|
Give it a Discord channel/thread id and an output directory; it exports the
|
|
channel, extracts the encoded download links (gofile.io and 0807.st), picks
|
|
the right downloader for each link, and keeps going when individual links
|
|
fail. Everything is recorded in state.json; failures also land in
|
|
failed_urls.txt so a re-run (or --urls-file) retries exactly those.
|
|
|
|
Already-downloaded URLs are recorded in a global store (default
|
|
~/.local/share/dlphelper/urls.json), so a link that was downloaded for one
|
|
channel is never re-downloaded for another - and re-runs of the same channel
|
|
skip everything finished (see also --assume-done for importing old lists).
|
|
|
|
Usage:
|
|
python3 dlp.py <channel_id> <out_dir> [options]
|
|
python3 dlp.py --subs subscriptions.json [options]
|
|
python3 dlp.py <channel_id> <out_dir> --urls-file failed_urls.txt
|
|
python3 dlp.py --assume-done old_urls.txt # import past downloads
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import getpass
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import random
|
|
import sys
|
|
import time
|
|
import zipfile
|
|
from itertools import chain
|
|
from pathlib import Path
|
|
|
|
from dlphelper import discord as dlp_discord
|
|
from dlphelper import gofile as dlph_gofile
|
|
from dlphelper import zero807 as dlph_zero807
|
|
from dlphelper.globalstore import GlobalStore, default_global_path
|
|
from dlphelper.state import State, append_manifest
|
|
|
|
log = logging.getLogger("dlp")
|
|
|
|
|
|
def sha256_of(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as fh:
|
|
for chunk in iter(lambda: fh.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def resolve_token(cli_value: str | None) -> str:
|
|
if cli_value:
|
|
return cli_value
|
|
if os.environ.get("DISCORD_TOKEN"):
|
|
return os.environ["DISCORD_TOKEN"]
|
|
log.info("No Discord token given ($DISCORD_TOKEN); prompting (input hidden).")
|
|
token = getpass.getpass("Discord token: ")
|
|
if not token:
|
|
raise SystemExit("No token provided, aborting.")
|
|
return token
|
|
|
|
|
|
def read_url_lines(path: Path) -> list[str]:
|
|
return [
|
|
line.strip()
|
|
for line in path.read_text(encoding="utf-8").splitlines()
|
|
if line.strip() and not line.strip().startswith("#")
|
|
]
|
|
|
|
|
|
def load_urls_file(path: Path) -> list[dlp_discord.FoundUrl]:
|
|
"""Load a plain URL list (e.g. failed_urls.txt); no attribution available."""
|
|
found: dict[str, dlp_discord.FoundUrl] = {}
|
|
for url in read_url_lines(path):
|
|
host = dlp_discord.classify(url)
|
|
if host is None:
|
|
log.warning("Unsupported host in %s: %s", path.name, url)
|
|
continue
|
|
found.setdefault(url, dlp_discord.FoundUrl(url=url, host=host))
|
|
return sorted(found.values(), key=lambda f: f.url)
|
|
|
|
|
|
def record_files(
|
|
state: State,
|
|
gstore: GlobalStore | None,
|
|
url: str,
|
|
files: list[dict],
|
|
source: dict,
|
|
) -> None:
|
|
state.mark(url, status="ok", error=None, files=files)
|
|
if gstore is not None:
|
|
gstore.mark(url, channel=source.get("channel"), host=source.get("host"),
|
|
file_count=len(files))
|
|
for f in files:
|
|
append_manifest(Path(source["manifest"]), {
|
|
"channel": source["channel"],
|
|
"url": url,
|
|
"host": source.get("host"),
|
|
"message_id": source.get("message_id"),
|
|
"timestamp": source.get("timestamp"),
|
|
"author": source.get("author"),
|
|
**f,
|
|
})
|
|
|
|
|
|
def extract_zips(files: list[dict]) -> None:
|
|
"""Optional: unzip each downloaded archive next to itself."""
|
|
for f in files:
|
|
path = Path(f["path"])
|
|
if path.suffix.lower() != ".zip" or not path.exists():
|
|
continue
|
|
target = path.with_suffix("")
|
|
if target.exists() and any(target.iterdir()):
|
|
log.info("Skip extracting %s (target exists)", path.name)
|
|
continue
|
|
log.info("Extracting %s -> %s", path.name, target.name)
|
|
try:
|
|
target.mkdir(parents=True, exist_ok=True)
|
|
with zipfile.ZipFile(path) as zf:
|
|
zf.extractall(target)
|
|
except (zipfile.BadZipFile, OSError) as exc:
|
|
log.warning("Extraction of %s failed: %s", path.name, exc)
|
|
|
|
|
|
def run_gofile(
|
|
urls: list[dlp_discord.FoundUrl],
|
|
args: argparse.Namespace,
|
|
state: State,
|
|
source_by_url: dict[str, dict],
|
|
is_done,
|
|
gstore: GlobalStore | None,
|
|
) -> None:
|
|
if not [f for f in urls if not is_done(f.url)]:
|
|
return # nothing pending: no guest account creation, no API calls
|
|
client = dlph_gofile.GoFileClient(password=args.password)
|
|
try:
|
|
client.ensure_token()
|
|
except dlph_gofile.GoFileError as exc:
|
|
log.error("%s - marking all pending gofile links as failed.", exc)
|
|
for f in urls:
|
|
if is_done(f.url):
|
|
continue
|
|
state.mark(f.url, status="failed", error=str(exc), files=[])
|
|
return
|
|
|
|
consecutive_failures = 0
|
|
for f in urls:
|
|
if is_done(f.url):
|
|
log.info("Skip (already downloaded): %s", f.url)
|
|
continue
|
|
source = source_by_url[f.url]
|
|
|
|
for attempt in range(1, args.retries + 1):
|
|
try:
|
|
files_spec = client.list_files(f.url)
|
|
if not files_spec:
|
|
raise dlph_gofile.GoFileError("no files found under URL")
|
|
files: list[dict] = []
|
|
for relpath, name, link in files_spec:
|
|
dest = Path(args.out_dir) / "downloads" / "gofile" / relpath / name
|
|
size = client.download_file(link, dest)
|
|
entry = {"path": str(dest), "size": size, "name": name}
|
|
if args.hash:
|
|
entry["sha256"] = sha256_of(dest)
|
|
files.append(entry)
|
|
log.info("OK %s (%d files)", f.url, len(files))
|
|
record_files(state, gstore, f.url, files, source)
|
|
if args.extract:
|
|
extract_zips(files)
|
|
consecutive_failures = 0
|
|
break
|
|
except dlph_gofile.GoFilePasswordRequired as exc:
|
|
# Retrying without a password never helps
|
|
log.error("PASSWORD REQUIRED %s: %s", f.url, exc)
|
|
state.mark(f.url, status="failed", error=str(exc), files=[])
|
|
break
|
|
except dlph_gofile.GoFileNotFound as exc:
|
|
# Dead links never come back - no retry
|
|
log.error("DEAD LINK %s: %s", f.url, exc)
|
|
state.mark(f.url, status="failed", error=str(exc), files=[])
|
|
break
|
|
except dlph_gofile.GoFileAuthError as exc:
|
|
# Token may have died mid-run (rate limit); re-auth once per attempt
|
|
log.warning("Auth error for %s: %s - re-authenticating", f.url, exc)
|
|
state.mark(f.url, error=str(exc), files=[])
|
|
client.reset()
|
|
try:
|
|
client.ensure_token()
|
|
except dlph_gofile.GoFileError as reauth_exc:
|
|
log.error("Re-authentication failed: %s", reauth_exc)
|
|
if attempt < args.retries:
|
|
time.sleep(args.retry_delay * attempt)
|
|
except dlph_gofile.GoFileError as exc:
|
|
log.warning("Attempt %d/%d failed for %s: %s",
|
|
attempt, args.retries, f.url, exc)
|
|
state.mark(f.url, error=str(exc), files=[])
|
|
if attempt < args.retries:
|
|
time.sleep(args.retry_delay * attempt)
|
|
except Exception as exc: # noqa: BLE001 - a bug must not stop the run
|
|
log.exception("Unexpected error for %s: %s", f.url, exc)
|
|
state.mark(f.url, error=f"unexpected: {exc}", files=[])
|
|
if attempt < args.retries:
|
|
time.sleep(args.retry_delay * attempt)
|
|
except KeyboardInterrupt:
|
|
state.mark(f.url, error="interrupted", files=[])
|
|
raise
|
|
else:
|
|
rec = state.records.get(f.url, {})
|
|
state.mark(f.url, status="failed", error=rec.get("error", "unknown"), files=[])
|
|
consecutive_failures += 1
|
|
log.error("FAILED %s after %d attempts", f.url, args.retries)
|
|
if consecutive_failures >= args.cooldown_failures:
|
|
log.warning(
|
|
"%d failures in a row - cooling down %ds before continuing "
|
|
"(rate limit?). The run does not stop.",
|
|
consecutive_failures, args.cooldown_seconds,
|
|
)
|
|
time.sleep(args.cooldown_seconds)
|
|
|
|
|
|
def run_0807(
|
|
urls: list[dlp_discord.FoundUrl],
|
|
args: argparse.Namespace,
|
|
state: State,
|
|
source_by_url: dict[str, dict],
|
|
is_done,
|
|
gstore: GlobalStore | None,
|
|
) -> None:
|
|
if not [f for f in urls if not is_done(f.url)]:
|
|
return
|
|
out_dir = Path(args.out_dir) / "downloads" / "0807"
|
|
consecutive_failures = 0
|
|
for f in urls:
|
|
if is_done(f.url):
|
|
log.info("Skip (already downloaded): %s", f.url)
|
|
continue
|
|
source = source_by_url[f.url]
|
|
|
|
for attempt in range(1, args.retries + 1):
|
|
try:
|
|
dest, size = dlph_zero807.download(f.url, out_dir)
|
|
entry: dict = {"path": str(dest), "size": size, "name": dest.name}
|
|
if args.hash:
|
|
entry["sha256"] = sha256_of(dest)
|
|
log.info("OK %s (%s)", f.url, format_size(size))
|
|
record_files(state, gstore, f.url, [entry], source)
|
|
if args.extract:
|
|
extract_zips([entry])
|
|
consecutive_failures = 0
|
|
break
|
|
except dlph_zero807.Zero807NotFound as exc:
|
|
# Dead links never come back - no retry
|
|
log.error("DEAD LINK %s: %s", f.url, exc)
|
|
state.mark(f.url, status="failed", error=str(exc), files=[])
|
|
break
|
|
except dlph_zero807.Zero807Error as exc:
|
|
log.warning("Attempt %d/%d failed for %s: %s",
|
|
attempt, args.retries, f.url, exc)
|
|
state.mark(f.url, error=str(exc), files=[])
|
|
if attempt < args.retries:
|
|
time.sleep(args.retry_delay * attempt)
|
|
except Exception as exc: # noqa: BLE001 - a bug must not stop the run
|
|
log.exception("Unexpected error for %s: %s", f.url, exc)
|
|
state.mark(f.url, error=f"unexpected: {exc}", files=[])
|
|
if attempt < args.retries:
|
|
time.sleep(args.retry_delay * attempt)
|
|
except KeyboardInterrupt:
|
|
state.mark(f.url, error="interrupted", files=[])
|
|
raise
|
|
else:
|
|
rec = state.records.get(f.url, {})
|
|
state.mark(f.url, status="failed", error=rec.get("error", "unknown"), files=[])
|
|
consecutive_failures += 1
|
|
log.error("FAILED %s after %d attempts", f.url, args.retries)
|
|
if consecutive_failures >= args.cooldown_failures:
|
|
log.warning(
|
|
"%d failures in a row - cooling down %ds before continuing "
|
|
"(HTTP 403 storm?). The run does not stop.",
|
|
consecutive_failures, args.cooldown_seconds,
|
|
)
|
|
time.sleep(args.cooldown_seconds)
|
|
|
|
|
|
def format_size(n: int) -> str:
|
|
size = float(n)
|
|
for unit in ("B", "KiB", "MiB", "GiB", "TiB"):
|
|
if size < 1024 or unit == "TiB":
|
|
return f"{size:.1f} {unit}"
|
|
size /= 1024
|
|
return f"{n} B"
|
|
|
|
|
|
def write_failed_urls(out_dir: Path, state: State) -> Path:
|
|
dest = out_dir / "failed_urls.txt"
|
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
|
failures = state.failures()
|
|
dest.write_text("\n".join(failures) + ("\n" if failures else ""), encoding="utf-8")
|
|
return dest
|
|
|
|
|
|
def mark_assume_done(path: Path, gstore: GlobalStore | None) -> int:
|
|
"""Record URLs as already downloaded (no files) so they are never fetched.
|
|
|
|
Manual import step for the pre-dlp.py era: point it at an old URL list you
|
|
know is fully downloaded. There is no verification - garbage in, skipped
|
|
forever.
|
|
"""
|
|
if gstore is None:
|
|
raise SystemExit("--assume-done needs the global store (drop --no-global)")
|
|
count = 0
|
|
for url in read_url_lines(path):
|
|
gstore.mark(url, channel=None, host=dlp_discord.classify(url), file_count=0,
|
|
note="assumed done via --assume-done")
|
|
count += 1
|
|
log.info("Marked %d URLs as done in the global store (%s)", count, gstore.path)
|
|
return count
|
|
|
|
|
|
def process_channel(
|
|
channel_id: str,
|
|
out_dir: Path,
|
|
args: argparse.Namespace,
|
|
token: str | None,
|
|
gstore: GlobalStore | None,
|
|
force_export: bool = False,
|
|
) -> int:
|
|
"""Run the full pipeline for one channel; returns the number of failures."""
|
|
state = State(out_dir / "state.json")
|
|
manifest = out_dir / "manifest.jsonl"
|
|
|
|
def is_done(url: str) -> bool:
|
|
return state.is_done(url) or (gstore is not None and gstore.is_done(url))
|
|
|
|
if args.urls_file:
|
|
found = load_urls_file(args.urls_file)
|
|
else:
|
|
if token is None:
|
|
raise SystemExit("internal error: no token resolved")
|
|
export_path = dlp_discord.export_channel(
|
|
channel_id, out_dir, token, force=force_export or args.fresh
|
|
)
|
|
found = dlp_discord.extract_urls(export_path)
|
|
|
|
if args.only:
|
|
found = [f for f in found if f.host == args.only]
|
|
|
|
buckets: dict[str, list[dlp_discord.FoundUrl]] = {"gofile": [], "0807": []}
|
|
for f in found:
|
|
if f.host in buckets:
|
|
buckets[f.host].append(f)
|
|
else: # classify() already filtered, but be defensive
|
|
log.warning("Skipping unsupported URL: %s", f.url)
|
|
state.mark(f.url, status="failed", error="unsupported host")
|
|
|
|
source_by_url: dict[str, dict] = {
|
|
f.url: {
|
|
"channel": channel_id,
|
|
"host": f.host,
|
|
"message_id": f.message_id,
|
|
"timestamp": f.timestamp,
|
|
"author": f.author,
|
|
"manifest": str(manifest),
|
|
}
|
|
for f in found
|
|
}
|
|
|
|
# --limit applies per host class to the *pending* URLs, so a test run
|
|
# downloads exactly N new things per host (done ones don't use up the quota)
|
|
if args.limit is not None:
|
|
for key, items in buckets.items():
|
|
pending = [f for f in items if not is_done(f.url)]
|
|
if args.shuffle:
|
|
random.shuffle(pending)
|
|
buckets[key] = pending[: args.limit]
|
|
|
|
pending = [f for f in found if not is_done(f.url)]
|
|
if args.limit is not None:
|
|
pending = [f for f in chain.from_iterable(buckets.values()) if not is_done(f.url)]
|
|
log.info(
|
|
"%s: %d URLs (%d gofile, %d 0807); %d done, %d to download",
|
|
channel_id, len(found),
|
|
len(buckets["gofile"]), len(buckets["0807"]),
|
|
len(found) - len(pending), len(pending),
|
|
)
|
|
|
|
if args.dry_run:
|
|
selected = {f.url for f in chain.from_iterable(buckets.values())}
|
|
for f in found:
|
|
if f.url in selected:
|
|
flag = "download"
|
|
elif is_done(f.url):
|
|
flag = "done"
|
|
else:
|
|
flag = "skip" # cut by --limit / --only
|
|
print(f"[{f.host:6}] [{flag:8}] {f.url}")
|
|
return 0
|
|
|
|
if buckets["gofile"]:
|
|
run_gofile(buckets["gofile"], args, state, source_by_url, is_done, gstore)
|
|
if buckets["0807"]:
|
|
run_0807(buckets["0807"], args, state, source_by_url, is_done, gstore)
|
|
|
|
failed_file = write_failed_urls(out_dir, state)
|
|
|
|
ok = [u for u, r in state.records.items() if r.get("status") == "ok"]
|
|
bad = state.failures()
|
|
dead = sum(1 for r in state.records.values() if "dead link" in (r.get("error") or ""))
|
|
total_files = sum(len(r.get("files") or []) for r in state.records.values())
|
|
total_size = sum(
|
|
(f_.get("size") or 0) for r in state.records.values() for f_ in (r.get("files") or [])
|
|
)
|
|
log.info("%s: %d ok (%d files, %s) / %d failed (%d dead links) / %d seen",
|
|
out_dir, len(ok), total_files, format_size(total_size),
|
|
len(bad), dead, len(state.records))
|
|
if bad:
|
|
log.info("Failed URLs written to %s - re-run this tool or use "
|
|
"python3 dlp.py <id> <out> --urls-file %s to retry them.",
|
|
failed_file, failed_file.name)
|
|
return len(bad)
|
|
|
|
|
|
def load_subscriptions(path: Path) -> list[dict]:
|
|
"""Load subscriptions.json: [{"name": ..., "channel_id": ..., "out_dir": ...}]."""
|
|
data = json.loads(path.read_text(encoding="utf-8"))
|
|
entries = data.get("subscriptions", []) if isinstance(data, dict) else data
|
|
if not isinstance(entries, list):
|
|
raise SystemExit(f"{path}: expected a list or {'{'}\"subscriptions\": [...]{'}'}")
|
|
subs = []
|
|
for entry in entries:
|
|
if not isinstance(entry, dict) or "channel_id" not in entry or "out_dir" not in entry:
|
|
raise SystemExit(
|
|
f"{path}: each subscription needs 'channel_id' and 'out_dir' "
|
|
f"(got: {entry})"
|
|
)
|
|
subs.append(entry)
|
|
return subs
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(
|
|
description="Collect downloads from Discord channels (gofile + 0807)."
|
|
)
|
|
parser.add_argument("channel_id", nargs="?", help="Discord channel or thread id")
|
|
parser.add_argument("out_dir", nargs="?", help="directory for downloads and state")
|
|
parser.add_argument("--subs", type=Path, metavar="FILE",
|
|
help="process a subscription list instead of one channel "
|
|
"(entries: channel_id, out_dir, optional name)")
|
|
parser.add_argument("--token", help="Discord token (else $DISCORD_TOKEN or prompt)")
|
|
parser.add_argument("--urls-file", type=Path,
|
|
help="skip the Discord export; read URLs from this file "
|
|
"(e.g. failed_urls.txt)")
|
|
parser.add_argument("--password", help="password for password-protected gofile links")
|
|
parser.add_argument("--retries", type=int, default=3, help="attempts per URL (default 3)")
|
|
parser.add_argument("--retry-delay", type=int, default=5,
|
|
help="base seconds between retries, multiplied per attempt (default 5)")
|
|
parser.add_argument("--cooldown-failures", type=int, default=3,
|
|
help="consecutive failures before a cooldown pause (default 3)")
|
|
parser.add_argument("--cooldown-seconds", type=int, default=60,
|
|
help="length of the cooldown pause (default 60)")
|
|
parser.add_argument("--only", choices=("gofile", "0807"),
|
|
help="process only one host class")
|
|
parser.add_argument("--limit", type=int, metavar="N",
|
|
help="download at most N pending URLs per host class "
|
|
"(test runs)")
|
|
parser.add_argument("--shuffle", action="store_true",
|
|
help="randomize which URLs a --limit picks")
|
|
parser.add_argument("--hash", action="store_true",
|
|
help="record sha256 of each downloaded file")
|
|
parser.add_argument("--extract", action="store_true",
|
|
help="unzip downloaded archives into per-archive folders")
|
|
parser.add_argument("--fresh", action="store_true",
|
|
help="re-export the Discord channel even if export.json exists "
|
|
"(subs mode does this by default)")
|
|
parser.add_argument("--no-fresh", action="store_true",
|
|
help="subs mode: reuse cached exports instead of re-exporting")
|
|
parser.add_argument("--dry-run", action="store_true",
|
|
help="export/extract/classify only, download nothing")
|
|
parser.add_argument("--global-state", type=Path, metavar="FILE",
|
|
help=f"cross-run URL store (default: {default_global_path()})")
|
|
parser.add_argument("--no-global", action="store_true",
|
|
help="disable the cross-run URL store for this run")
|
|
parser.add_argument("--assume-done", type=Path, metavar="FILE",
|
|
help="mark every URL in the file as already downloaded, "
|
|
"then exit (manual import of old download lists)")
|
|
args = parser.parse_args()
|
|
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s %(levelname)-7s %(message)s",
|
|
datefmt="%H:%M:%S",
|
|
)
|
|
|
|
gstore = None
|
|
if not args.no_global:
|
|
gstore = GlobalStore(args.global_state or default_global_path())
|
|
|
|
if args.assume_done:
|
|
if args.channel_id or args.subs:
|
|
parser.error("--assume-done cannot be combined with a channel or --subs")
|
|
mark_assume_done(args.assume_done, gstore)
|
|
return
|
|
|
|
if args.subs:
|
|
if args.channel_id or args.out_dir or args.urls_file:
|
|
parser.error("--subs cannot be combined with channel_id/out_dir/--urls-file")
|
|
subs = load_subscriptions(args.subs)
|
|
token = resolve_token(args.token)
|
|
total_failed = 0
|
|
for entry in subs:
|
|
out_dir = Path(entry["out_dir"])
|
|
log.info("=== subscription: %s ===", entry.get("name") or entry["channel_id"])
|
|
total_failed += process_channel(
|
|
entry["channel_id"], out_dir, args, token, gstore,
|
|
force_export=not args.no_fresh, # subs mode re-exports by default
|
|
)
|
|
if total_failed:
|
|
log.warning("%d URL(s) failed in total across subscriptions.", total_failed)
|
|
sys.exit(1)
|
|
return
|
|
|
|
if not args.channel_id or not args.out_dir:
|
|
parser.error("channel_id and out_dir are required (or use --subs/--assume-done)")
|
|
|
|
# A token is only needed when we will export a channel (--urls-file skips that)
|
|
token = None if args.urls_file else resolve_token(args.token)
|
|
failures = process_channel(args.channel_id, Path(args.out_dir), args, token, gstore)
|
|
if failures:
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
main()
|
|
except KeyboardInterrupt:
|
|
log.warning("Interrupted - state.json is saved; re-run to resume.")
|
|
sys.exit(130)
|