pull-messages-from-discord/script.py

# /// script
# requires-python = ">=3.11"
# dependencies = ["httpx"]
# ///
"""Pull messages from a Discord channel for the last N hours and save them as JSON."""

import argparse
import json
import os
import re
import sys
import time
from datetime import datetime, timedelta, timezone
from pathlib import Path

import httpx

API = "https://discord.com/api/v10"
DISCORD_EPOCH_MS = 1420070400000
SCRIPT_DIR = Path(__file__).resolve().parent
# 0 text, 5 announcement
TEXT_CHANNEL_TYPES = {0, 5}


def load_token() -> str:
    token = os.environ.get("HM_DISCORD_BOT_TOKEN")
    if token:
        return token
    for name in (".env", "discord-bot.txt"):
        path = SCRIPT_DIR / name
        if not path.is_file():
            continue
        for line in path.read_text().splitlines():
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            key, _, value = line.partition("=")
            if key.strip() == "HM_DISCORD_BOT_TOKEN":
                return value.strip().strip("'\"")
    sys.exit("error: HM_DISCORD_BOT_TOKEN not found in environment, .env, or discord-bot.txt")


def parse_hours(value: str) -> int:
    m = re.fullmatch(r"(\d+)h?", value.strip().lower())
    if not m or int(m.group(1)) == 0:
        raise argparse.ArgumentTypeError(f"invalid time {value!r}; use hours, e.g. 24 or 24h")
    return int(m.group(1))


def datetime_to_snowflake(dt: datetime) -> int:
    return (int(dt.timestamp() * 1000) - DISCORD_EPOCH_MS) << 22


def api_get(client: httpx.Client, path: str, params: dict | None = None):
    while True:
        resp = client.get(f"{API}{path}", params=params)
        if resp.status_code == 429:
            time.sleep(float(resp.json().get("retry_after", 1)) + 0.1)
            continue
        if resp.status_code == 401:
            sys.exit("error: 401 unauthorized - bot token invalid")
        if resp.status_code == 403:
            sys.exit(f"error: 403 forbidden on {path} - bot needs View Channel + Read Message History")
        if resp.status_code >= 400:
            sys.exit(f"error: {resp.status_code} on {path}: {resp.text}")
        return resp.json()


def find_channel(client: httpx.Client, name: str, guild_filter: str | None) -> tuple[dict, dict]:
    matches = []
    for guild in api_get(client, "/users/@me/guilds"):
        if guild_filter and guild_filter not in (guild["id"], guild["name"]):
            continue
        for ch in api_get(client, f"/guilds/{guild['id']}/channels"):
            if ch["type"] in TEXT_CHANNEL_TYPES and name in (ch["name"], ch["id"]):
                matches.append((guild, ch))
    if not matches:
        sys.exit(f"error: no text channel named {name!r} visible to bot")
    if len(matches) > 1:
        options = "\n".join(
            f"  {ch['id']}  guild={g['name']!r}  parent_id={ch.get('parent_id')}" for g, ch in matches
        )
        sys.exit(
            f"error: {len(matches)} channels match {name!r}:\n{options}\n"
            "Pass a channel ID to --channel (or --guild to narrow by server)."
        )
    return matches[0]


def fetch_messages(client: httpx.Client, channel_id: str, since: datetime) -> list[dict]:
    messages: list[dict] = []
    after = str(datetime_to_snowflake(since))
    while True:
        page = api_get(client, f"/channels/{channel_id}/messages", {"limit": 100, "after": after})
        if not page:
            break
        messages.extend(page)
        after = max(page, key=lambda m: int(m["id"]))["id"]
        if len(page) < 100:
            break
    messages.sort(key=lambda m: int(m["id"]))
    return messages


def slim(msg: dict) -> dict:
    author = msg.get("author", {})
    return {
        "id": msg["id"],
        "timestamp": msg["timestamp"],
        "edited_timestamp": msg.get("edited_timestamp"),
        "author": {
            "id": author.get("id"),
            "username": author.get("username"),
            "bot": author.get("bot", False),
        },
        "content": msg.get("content", ""),
        "embeds": msg.get("embeds", []),
        "attachments": [
            {"filename": a.get("filename"), "url": a.get("url")} for a in msg.get("attachments", [])
        ],
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("-c", "--channel", required=True, help="channel name, e.g. homelab-alerts")
    parser.add_argument("-t", "--time", type=parse_hours, default=24, help="hours back, e.g. 24 or 24h (default 24)")
    parser.add_argument("-g", "--guild", help="guild name or ID (only if channel name is ambiguous)")
    parser.add_argument("-o", "--output", type=Path, help="output file (default: next to script)")
    args = parser.parse_args()

    channel_name = args.channel.lstrip("#")
    now = datetime.now(timezone.utc)
    since = now - timedelta(hours=args.time)

    headers = {"Authorization": f"Bot {load_token()}", "User-Agent": "DiscordBot (pull-messages, 1.0)"}
    with httpx.Client(headers=headers, timeout=30) as client:
        guild, channel = find_channel(client, channel_name, args.guild)
        raw = fetch_messages(client, channel["id"], since)

    result = {
        "channel": channel_name,
        "channel_id": channel["id"],
        "guild": guild["name"],
        "since": since.isoformat(),
        "fetched_at": now.isoformat(),
        "hours": args.time,
        "count": len(raw),
        "messages": [slim(m) for m in raw],
    }
    out = args.output or SCRIPT_DIR / f"discord_{channel_name}_{now:%Y%m%dT%H%M%SZ}.json"
    out.write_text(json.dumps(result, indent=2, ensure_ascii=False) + "\n")
    print(f"Saved {len(raw)} messages -> {out}")


if __name__ == "__main__":
    main()