Warhammer_DLP_Helper/src/discord.rs

//! DCE docker export (JSON) + base64/plain URL extraction — Rust port of
//! `dlphelper/discord.py`.
//!
//! 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.

use std::collections::HashMap;
use std::fs;
use std::io;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::sync::LazyLock;

use anyhow::{Context, Result, bail};
use base64::Engine as _;
use base64::engine::general_purpose::STANDARD;
use regex::Regex;
use serde_json::Value;
use url::Url;

pub const DCE_IMAGE: &str = "tyrrrz/discordchatexporter:stable";

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

/// Host key ("gofile" | "0807") -> host names accepted for it.
pub const SUPPORTED_HOSTS: &[(&str, &[&str])] = &[
    ("gofile", &["gofile.io", "www.gofile.io"]),
    ("0807", &["0807.st", "www.0807.st"]),
];

/// One download URL found in the channel, with source attribution.
#[derive(Debug, Clone, PartialEq)]
pub struct FoundUrl {
    pub url: String,
    /// "gofile" | "0807" (set after classification)
    pub host: &'static str,
    pub message_id: Option<String>,
    pub timestamp: Option<String>,
    pub author: Option<String>,
    /// found via base64 decode vs. plain link
    pub encoded: bool,
}

/// 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.
pub fn export_channel(
    channel_id: &str,
    out_dir: &Path,
    token: &str,
    force: bool,
) -> Result<PathBuf> {
    let dest = out_dir.join("export.json");
    if dest.exists() && !force {
        log::info!(
            "Reusing existing export {} (pass --fresh to re-export)",
            dest.display()
        );
        return Ok(dest);
    }

    fs::create_dir_all(out_dir)?;
    // Python Path.resolve() on an existing dir == canonicalize here (the
    // create_dir_all above guarantees it exists).
    let resolved = fs::canonicalize(out_dir)?;
    let (uid, gid) = read_uid_gid()?;

    log::info!(
        "Exporting channel {} via DiscordChatExporter ...",
        channel_id
    );
    // stdout/stderr are inherited, like the python subprocess.run(check=True).
    let status = Command::new("docker")
        .args([
            "run",
            "--rm",
            "--user",
            &format!("{uid}:{gid}"),
            "-v",
            &format!("{}:/out", resolved.display()),
            "-e",
            "DISCORD_TOKEN",
            "-e",
            "TOKEN",
        ])
        .arg(DCE_IMAGE)
        .args([
            "export",
            "-c",
            channel_id,
            "-f",
            "Json",
            "-o",
            "/out/export.json",
        ])
        // The token only travels via the environment, never on the argv.
        .env("DISCORD_TOKEN", token)
        .env("TOKEN", token)
        .status();
    match status {
        Err(e) if e.kind() == io::ErrorKind::NotFound => {
            bail!("docker not found - the Discord export needs the docker image {DCE_IMAGE}")
        }
        Err(e) => return Err(e).context("failed to spawn docker"),
        Ok(s) if !s.success() => bail!(
            "DiscordChatExporter failed with exit code {} (check the token and the channel id)",
            s.code().unwrap_or(-1)
        ),
        Ok(_) => {}
    }
    if !dest.exists() {
        bail!("DiscordChatExporter produced no output file");
    }
    Ok(dest)
}

/// Real uid/gid of this process. Linux only: parsed from /proc/self/status
/// because std has no getuid (the first number on the Uid:/Gid: line is the
/// real id, matching python's os.getuid()/os.getgid()).
fn read_uid_gid() -> Result<(u32, u32)> {
    let status = fs::read_to_string("/proc/self/status")
        .context("cannot read /proc/self/status to determine uid/gid for docker --user")?;
    let uid = parse_proc_id(&status, "Uid:")?;
    let gid = parse_proc_id(&status, "Gid:")?;
    Ok((uid, gid))
}

fn parse_proc_id(status_text: &str, key: &str) -> Result<u32> {
    for line in status_text.lines() {
        if let Some(rest) = line.strip_prefix(key) {
            let first = rest
                .split_whitespace()
                .next()
                .with_context(|| format!("empty {key} line in /proc/self/status"))?;
            return first.parse().with_context(|| {
                format!("cannot parse {key} value {first:?} in /proc/self/status")
            });
        }
    }
    bail!("{key} line not found in /proc/self/status")
}

/// Return the host key ('gofile' | '0807') for a URL, else None.
///
/// Mirrors python `urlparse(url).hostname.lower()`: unparseable URLs and URLs
/// without a hostname yield None. The url crate normalizes the host to
/// lowercase, same as urlparse does.
pub fn classify(url: &str) -> Option<&'static str> {
    let parsed = Url::parse(url).ok()?;
    let hostname = parsed.host_str()?;
    SUPPORTED_HOSTS
        .iter()
        .find(|(_, hosts)| hosts.contains(&hostname))
        .map(|(key, _)| *key)
}

/// Decode one base64 candidate into the URL it contains, or None.
///
/// Python quirk kept: "_" is remapped to "/" and "-" to "+" (in that order),
/// trailing "=" stripped, then padded back to a multiple of 4 before a strict
/// standard-alphabet decode and a strict UTF-8 check. One encoded blob can
/// contain the URL plus padding/newlines: whitespace is trimmed and only the
/// first line is returned.
fn decode_b64(candidate: &str) -> Option<String> {
    let mut s: String = candidate
        .chars()
        .map(|c| match c {
            '_' => '/',
            '-' => '+',
            other => other,
        })
        .collect();
    while s.ends_with('=') {
        s.pop();
    }
    for _ in 0..(4 - s.len() % 4) % 4 {
        s.push('=');
    }
    let decoded = STANDARD.decode(s.as_bytes()).ok()?;
    let text = String::from_utf8(decoded).ok()?;
    // strip whitespace, keep the first line ("" -> None)
    text.trim()
        .lines()
        .next()
        .filter(|l| !l.is_empty())
        .map(String::from)
}

/// Python `str(x or "")` for a JSON value: falsy values (null, false, 0, "",
/// empty array/object) become the empty string, strings pass through, and any
/// other value is stringified. DCE content/id fields are always strings; the
/// other shapes only exist for parity with the python code.
fn python_str_or_empty(value: Option<&Value>) -> String {
    let falsy = match value {
        None | Some(Value::Null) => true,
        Some(Value::Bool(b)) => !b,
        Some(Value::Number(n)) => n.as_f64().is_some_and(|f| f == 0.0),
        Some(Value::String(s)) => s.is_empty(),
        Some(Value::Array(a)) => a.is_empty(),
        Some(Value::Object(o)) => o.is_empty(),
    };
    if falsy {
        return String::new();
    }
    match value.unwrap() {
        Value::String(s) => s.clone(),
        Value::Number(n) => n.to_string(),
        Value::Bool(_) => "True".to_string(), // python str(True)
        other => other.to_string(),
    }
}

/// Extract supported URLs from a DCE JSON export, deduplicated.
///
/// The root may be a list of messages or an object with a "messages" array;
/// anything else yields no URLs. The first occurrence of a URL wins
/// (python dict.setdefault semantics).
pub fn extract_urls(export_path: &Path) -> Result<Vec<FoundUrl>> {
    let text = fs::read_to_string(export_path)?;
    let data: Value = serde_json::from_str(&text)?;
    let empty: Vec<Value> = Vec::new();
    let messages: &[Value] = match &data {
        Value::Array(list) => list,
        Value::Object(obj) => obj
            .get("messages")
            .and_then(Value::as_array)
            .map(|a| a.as_slice())
            .unwrap_or(&empty),
        _ => &empty,
    };

    let mut found: HashMap<String, FoundUrl> = HashMap::new();
    let mut stat_b64 = 0u64;
    let mut stat_plain = 0u64;
    let mut stat_unsupported = 0u64;

    for msg in messages {
        let content = python_str_or_empty(msg.get("content"));
        if content.is_empty() {
            continue;
        }
        let author = msg
            .get("author")
            .and_then(|a| a.get("name"))
            .and_then(Value::as_str)
            .map(str::to_owned);
        let message_id = {
            let s = python_str_or_empty(msg.get("id"));
            (!s.is_empty()).then_some(s)
        };
        let timestamp = msg
            .get("timestamp")
            .and_then(Value::as_str)
            .map(str::to_owned);

        for m in B64_RE.find_iter(&content) {
            stat_b64 += 1;
            let Some(decoded) = decode_b64(m.as_str()) else {
                continue;
            };
            let Some(host) = classify(&decoded) else {
                stat_unsupported += 1;
                continue;
            };
            found.entry(decoded.clone()).or_insert(FoundUrl {
                url: decoded,
                host,
                message_id: message_id.clone(),
                timestamp: timestamp.clone(),
                author: author.clone(),
                encoded: true,
            });
        }

        for m in URL_RE.find_iter(&content) {
            let url = m
                .as_str()
                .trim_end_matches(['.', ',', ';', ')'])
                .to_string();
            let Some(host) = classify(&url) else {
                continue;
            };
            stat_plain += 1;
            found.entry(url.clone()).or_insert(FoundUrl {
                url,
                host,
                message_id: message_id.clone(),
                timestamp: timestamp.clone(),
                author: author.clone(),
                encoded: false,
            });
        }
    }

    log::info!(
        "Extraction: {} base64 candidates, {} plain links, {} unsupported/undecodable, {} unique supported URLs",
        stat_b64,
        stat_plain,
        stat_unsupported,
        found.len()
    );
    let mut out: Vec<FoundUrl> = found.into_values().collect();
    out.sort_by(|a, b| a.url.cmp(&b.url));
    Ok(out)
}

#[cfg(test)]
mod tests {
    use super::*;
    use base64::engine::general_purpose::URL_SAFE_NO_PAD;
    use serde_json::json;
    use tempfile::tempdir;

    fn b64(s: &str) -> String {
        STANDARD.encode(s.as_bytes())
    }

    /// The two-message fixture from tests/test_extract.py.
    fn fixture_messages() -> Value {
        json!([
            {
                "id": "1",
                "timestamp": "2026-01-01T00:00:00",
                "author": {"name": "poster"},
                "content": format!("chaos stuff: {}", b64("https://gofile.io/d/abc123")),
            },
            {
                "id": "2",
                "timestamp": "2026-01-02T00:00:00",
                "author": {"name": "poster"},
                "content": format!(
                    "0807 zip: {}\ndupe of 1: {}\nplain: https://0807.st/plainLink.zip and \
                     https://discord.com/channels/1/2 (unsupported)\njunk: {} and \
                     not-our-domain base64: aHR0cHM6Ly9leGFtcGxlLm9yZw==",
                    b64("https://0807.st/6saCNie.zip"),
                    b64("https://gofile.io/d/abc123"),
                    b64("https://pixeldrain.com/u/xyz"),
                ),
            },
        ])
    }

    fn write_export(dir: &Path, root: Value) -> PathBuf {
        let path = dir.join("export.json");
        fs::write(&path, serde_json::to_string(&root).unwrap()).unwrap();
        path
    }

    #[test]
    fn extract_urls_matches_python_fixture() {
        let dir = tempdir().unwrap();
        let export = write_export(dir.path(), json!({"messages": fixture_messages()}));
        let found = extract_urls(&export).unwrap();

        let urls: Vec<&str> = found.iter().map(|f| f.url.as_str()).collect();
        assert_eq!(
            urls,
            vec![
                "https://0807.st/6saCNie.zip",
                "https://0807.st/plainLink.zip",
                "https://gofile.io/d/abc123",
            ]
        );

        let g = found
            .iter()
            .find(|f| f.url == "https://gofile.io/d/abc123")
            .unwrap();
        assert_eq!(g.host, "gofile");
        assert_eq!(g.message_id.as_deref(), Some("1"));
        assert_eq!(g.timestamp.as_deref(), Some("2026-01-01T00:00:00"));
        assert_eq!(g.author.as_deref(), Some("poster"));
        assert!(g.encoded);

        let z = found
            .iter()
            .find(|f| f.url == "https://0807.st/6saCNie.zip")
            .unwrap();
        assert_eq!(z.host, "0807");
        assert_eq!(z.message_id.as_deref(), Some("2"));
        assert!(z.encoded);

        let p = found
            .iter()
            .find(|f| f.url == "https://0807.st/plainLink.zip")
            .unwrap();
        assert_eq!(p.host, "0807");
        assert_eq!(p.message_id.as_deref(), Some("2"));
        assert_eq!(p.author.as_deref(), Some("poster"));
        assert!(!p.encoded);
    }

    #[test]
    fn extract_urls_accepts_bare_list_root() {
        let dir = tempdir().unwrap();
        let export = write_export(dir.path(), fixture_messages());
        let found = extract_urls(&export).unwrap();
        assert_eq!(found.len(), 3);
    }

    #[test]
    fn extract_urls_empty_and_malformed_roots_yield_nothing() {
        let dir = tempdir().unwrap();
        assert!(
            extract_urls(&write_export(dir.path(), json!({})))
                .unwrap()
                .is_empty()
        );
        assert!(
            extract_urls(&write_export(dir.path(), json!({"other": []})))
                .unwrap()
                .is_empty()
        );
        // message without content is skipped
        let export = write_export(dir.path(), json!([{"id": "7"}]));
        assert!(extract_urls(&export).unwrap().is_empty());
    }

    #[test]
    fn classify_matches_python_hostnames() {
        assert_eq!(classify("https://0807.st/x.zip"), Some("0807"));
        assert_eq!(classify("https://www.0807.st/x.zip"), Some("0807"));
        assert_eq!(classify("https://gofile.io/d/x"), Some("gofile"));
        assert_eq!(classify("https://www.gofile.io/d/x"), Some("gofile"));
        assert_eq!(classify("https://example.org/x"), None);
        // url crate lowercases hosts, like urlparse().hostname.lower()
        assert_eq!(classify("https://GOFILE.IO/d/x"), Some("gofile"));
        // edge cases: no scheme, empty string, bare host
        assert_eq!(classify("0807.st/x.zip"), None);
        assert_eq!(classify(""), None);
        assert_eq!(classify("https://"), None);
    }

    #[test]
    fn decode_b64_handles_padding_and_rstrip() {
        // candidate as matched by B64_RE carries no '=': padding is added
        let unpadded = b64("https://gofile.io/d/abc123")
            .trim_end_matches('=')
            .to_string();
        assert_eq!(
            decode_b64(&unpadded),
            Some("https://gofile.io/d/abc123".to_string())
        );
        // trailing '=' is stripped before re-padding
        assert_eq!(
            decode_b64("aHR0cHM6Ly9leGFtcGxlLm9yZw=="),
            Some("https://example.org".to_string())
        );
    }

    #[test]
    fn decode_b64_remaps_urlsafe_alphabet() {
        // U+13C0 encodes to 4Y+A (standard) / 4Y-A (urlsafe); the python
        // quirk remaps '-' back to '+' before decoding. ('_' -> '/' can never
        // occur in a candidate that decodes to valid UTF-8, so '-' covers the
        // remap path.)
        let candidate = URL_SAFE_NO_PAD.encode("\u{13C0}");
        assert!(candidate.contains('-'), "expected '-' in {candidate}");
        assert_eq!(decode_b64(&candidate), Some("\u{13C0}".to_string()));
    }

    #[test]
    fn decode_b64_keeps_first_line_after_whitespace() {
        let candidate = STANDARD.encode("https://gofile.io/d/abc123\nnoise  ");
        assert_eq!(
            decode_b64(&candidate),
            Some("https://gofile.io/d/abc123".to_string())
        );
    }

    #[test]
    fn decode_b64_rejects_invalid_utf8_and_garbage() {
        // decodes to \xff\xfe — not valid UTF-8
        assert_eq!(decode_b64("//8+"), None);
        assert_eq!(decode_b64("aHR0cHM6Ly["), None);
        assert_eq!(decode_b64(""), None);
    }

    #[test]
    fn parse_proc_id_reads_first_number() {
        let text = "Name:\tx\nUid:\t1000\t1000\t1000\t1000\nGid:\t989\t989\t989\t989\n";
        assert_eq!(parse_proc_id(text, "Uid:").unwrap(), 1000);
        assert_eq!(parse_proc_id(text, "Gid:").unwrap(), 989);
        assert!(parse_proc_id("Name:\tx\n", "Uid:").is_err());
    }

    #[test]
    fn read_uid_gid_matches_proc_self() {
        // Linux-only binary (docker export); /proc/self/status must exist.
        let (uid, gid) = read_uid_gid().unwrap();
        let text = fs::read_to_string("/proc/self/status").unwrap();
        let want_uid = parse_proc_id(&text, "Uid:").unwrap();
        let want_gid = parse_proc_id(&text, "Gid:").unwrap();
        assert_eq!((uid, gid), (want_uid, want_gid));
    }

    #[test]
    fn export_channel_reuses_existing_export_without_docker() {
        let dir = tempdir().unwrap();
        let dest = dir.path().join("export.json");
        fs::write(&dest, "[]").unwrap();
        // force=false must short-circuit before any docker invocation; the
        // token must not end up in the error/log path either.
        let got = export_channel("123", dir.path(), "s3cret-token", false).unwrap();
        assert_eq!(got, dest);
    }
}