Warhammer_DLP_Helper/dlphelper/discord.py

"""Export a Discord channel/thread with DiscordChatExporter and extract URLs.

The channels this tool targets encode their download links as base64 inside
message content, and occasionally as plain links. We export the channel as
JSON (richer metadata than the HTML export: message id, timestamp, author),
decode every base64 candidate, and keep only URLs for supported hosts.
"""

from __future__ import annotations

import base64
import json
import logging
import os
import re
import subprocess
from dataclasses import dataclass
from pathlib import Path
from urllib.parse import urlparse

log = logging.getLogger(__name__)

DCE_IMAGE = "tyrrrz/discordchatexporter:stable"

# The channels encode links as base64 that always starts with this prefix
# (https://...); grab.sh established this regex, keep it compatible.
B64_RE = re.compile(r"aHR0cHM6Ly[89][A-Za-z0-9+/_-]*")
URL_RE = re.compile(r"https://[^\s<>\"'\)\]\}]+")

SUPPORTED_HOSTS = {
    "gofile": ("gofile.io", "www.gofile.io"),
    "0807": ("0807.st", "www.0807.st"),
}


@dataclass(frozen=True)
class FoundUrl:
    """One download URL found in the channel, with source attribution."""

    url: str
    host: str  # "gofile" | "0807" (set after classification)
    message_id: str | None = None
    timestamp: str | None = None
    author: str | None = None
    encoded: bool = False  # found via base64 decode vs. plain link


def export_channel(channel_id: str, out_dir: Path, token: str, force: bool = False) -> Path:
    """Export the channel to out_dir/export.json using the DCE docker image.

    The token is passed to the container via the DISCORD_TOKEN env var (DCE
    reads it natively) so it never appears on the docker command line. TOKEN
    is set as well for DCE versions that only know that name.
    """
    dest = out_dir / "export.json"
    if dest.exists() and not force:
        log.info("Reusing existing export %s (pass --fresh to re-export)", dest)
        return dest

    out_dir.mkdir(parents=True, exist_ok=True)
    env = dict(os.environ)
    env["DISCORD_TOKEN"] = token
    env["TOKEN"] = token
    cmd = [
        "docker", "run", "--rm",
        "--user", f"{os.getuid()}:{os.getgid()}",
        "-v", f"{out_dir.resolve()}:/out",
        "-e", "DISCORD_TOKEN",
        "-e", "TOKEN",
        DCE_IMAGE,
        "export", "-c", channel_id, "-f", "Json", "-o", "/out/export.json",
    ]
    log.info("Exporting channel %s via DiscordChatExporter ...", channel_id)
    try:
        subprocess.run(cmd, env=env, check=True)
    except FileNotFoundError as exc:
        raise RuntimeError(
            "docker not found - the Discord export needs the docker image "
            f"{DCE_IMAGE}"
        ) from exc
    except subprocess.CalledProcessError as exc:
        raise RuntimeError(
            f"DiscordChatExporter failed with exit code {exc.returncode} "
            "(check the token and the channel id)"
        ) from exc
    if not dest.exists():
        raise RuntimeError("DiscordChatExporter produced no output file")
    return dest


def classify(url: str) -> str | None:
    """Return the host key ('gofile' | '0807') for a URL, else None."""
    try:
        hostname = (urlparse(url).hostname or "").lower()
    except ValueError:
        return None
    for key, hosts in SUPPORTED_HOSTS.items():
        if hostname in hosts:
            return key
    return None


def _decode_b64(candidate: str) -> str | None:
    s = candidate.replace("_", "/").replace("-", "+")
    s = s.rstrip("=")
    s += "=" * ((4 - len(s) % 4) % 4)
    try:
        decoded = base64.b64decode(s, validate=False).decode("utf-8", "strict")
    except Exception:  # noqa: BLE001 - malformed candidate, just skip it
        return None
    # One encoded blob can contain the URL plus padding/newlines
    stripped = decoded.strip()
    line = stripped.splitlines()[0] if stripped else ""
    return line or None


def extract_urls(export_path: Path) -> list[FoundUrl]:
    """Extract supported URLs from a DCE JSON export, deduplicated."""
    data = json.loads(export_path.read_text(encoding="utf-8"))
    messages = data if isinstance(data, list) else data.get("messages", [])

    found: dict[str, FoundUrl] = {}
    stats = {"b64": 0, "plain": 0, "unsupported": 0}

    for msg in messages:
        content = str(msg.get("content") or "")
        if not content:
            continue
        author = (msg.get("author") or {}).get("name")
        attrib = {
            "message_id": str(msg.get("id") or "") or None,
            "timestamp": msg.get("timestamp"),
            "author": author,
        }

        for match in B64_RE.finditer(content):
            stats["b64"] += 1
            decoded = _decode_b64(match.group(0))
            if not decoded:
                continue
            host = classify(decoded)
            if host is None:
                stats["unsupported"] += 1
                continue
            found.setdefault(decoded, FoundUrl(url=decoded, host=host, encoded=True, **attrib))

        for match in URL_RE.finditer(content):
            url = match.group(0).rstrip(".,;)")
            host = classify(url)
            if host is None:
                continue
            stats["plain"] += 1
            found.setdefault(url, FoundUrl(url=url, host=host, encoded=False, **attrib))

    log.info(
        "Extraction: %d base64 candidates, %d plain links, %d unsupported/undecodable, "
        "%d unique supported URLs",
        stats["b64"], stats["plain"], stats["unsupported"], len(found),
    )
    return sorted(found.values(), key=lambda f: f.url)