TheOrb/archive_bot/discord_api.py

from __future__ import annotations

import asyncio
import json
import secrets
import sys
import threading
import urllib.error
import urllib.request
from typing import Any

from .core import DISCORD_API


class DiscordGatewayManager:
    def __init__(self, token: str) -> None:
        self.token = token
        self.thread: threading.Thread | None = None
        self.loop: asyncio.AbstractEventLoop | None = None
        self.client: Any = None
        self.ready = threading.Event()
        self._disconnecting = threading.Event()

    def start(self) -> None:
        if self.thread is not None:
            return
        self.thread = threading.Thread(target=self._run, name="discord-gateway", daemon=True)
        self.thread.start()

    def stop(self) -> None:
        self._disconnecting.set()
        if self.loop is not None and self.client is not None:
            asyncio.run_coroutine_threadsafe(self.client.close(), self.loop)
        if self.thread is not None:
            self.thread.join(timeout=15)

    def _run(self) -> None:
        try:
            import discord
        except ImportError:
            print(
                "discord.py is not installed. Install requirements.txt to keep the bot online.",
                file=sys.stderr,
                flush=True,
            )
            self.ready.set()
            return

        class GatewayClient(discord.Client):
            def __init__(self, manager: DiscordGatewayManager) -> None:
                intents = discord.Intents.default()
                intents.guilds = True
                super().__init__(intents=intents)
                self.manager = manager

            async def on_ready(self) -> None:
                await self.change_presence(status=discord.Status.online)
                user = self.user
                name = user.name if user is not None else "unknown"
                bot_id = user.id if user is not None else "unknown"
                print(f"Discord gateway connected as {name} ({bot_id})", flush=True)
                self.manager.ready.set()

            async def on_disconnect(self) -> None:
                self.manager.ready.clear()

        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)
        self.client = GatewayClient(self)

        try:
            self.loop.run_until_complete(self.client.start(self.token))
        except Exception as exc:
            if not self._disconnecting.is_set():
                print(f"Discord gateway stopped: {exc}", file=sys.stderr, flush=True)
        finally:
            self.ready.set()
            try:
                pending = asyncio.all_tasks(self.loop)
                for task in pending:
                    task.cancel()
                if pending:
                    self.loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
            finally:
                self.loop.close()


def discord_request(
    method: str,
    token: str,
    path: str,
    payload: dict[str, Any] | None = None,
) -> dict[str, Any]:
    body = None
    headers = {
        "Authorization": f"Bot {token}",
        "User-Agent": "ArchiveStatusBot/1.0",
    }

    if payload is not None:
        body = json.dumps(payload).encode("utf-8")
        headers["Content-Type"] = "application/json"

    request = urllib.request.Request(
        f"{DISCORD_API}{path}",
        data=body,
        headers=headers,
        method=method,
    )

    try:
        with urllib.request.urlopen(request, timeout=20) as response:
            data = response.read()
            if not data:
                return {}
            return json.loads(data.decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="ignore")
        raise RuntimeError(f"Discord API {method} {path} failed: {exc.code} {detail}") from exc


def discord_multipart_request(
    method: str,
    token: str,
    path: str,
    payload: dict[str, Any],
    files: list[tuple[str, str, str, bytes]],
) -> dict[str, Any]:
    boundary = f"----ArchiveBot{secrets.token_hex(16)}"
    body = bytearray()

    def add_part(name: str, content: bytes, content_type: str, filename: str | None = None) -> None:
        body.extend(f"--{boundary}\r\n".encode("ascii"))
        disposition = f'Content-Disposition: form-data; name="{name}"'
        if filename is not None:
            disposition += f'; filename="{filename}"'
        body.extend(f"{disposition}\r\n".encode("utf-8"))
        body.extend(f"Content-Type: {content_type}\r\n\r\n".encode("ascii"))
        body.extend(content)
        body.extend(b"\r\n")

    add_part("payload_json", json.dumps(payload).encode("utf-8"), "application/json")
    for field_name, filename, content_type, content in files:
        add_part(field_name, content, content_type, filename)
    body.extend(f"--{boundary}--\r\n".encode("ascii"))

    request = urllib.request.Request(
        f"{DISCORD_API}{path}",
        data=bytes(body),
        headers={
            "Authorization": f"Bot {token}",
            "User-Agent": "ArchiveStatusBot/1.0",
            "Content-Type": f"multipart/form-data; boundary={boundary}",
        },
        method=method,
    )

    try:
        with urllib.request.urlopen(request, timeout=30) as response:
            data = response.read()
            if not data:
                return {}
            return json.loads(data.decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="ignore")
        raise RuntimeError(f"Discord API {method} {path} failed: {exc.code} {detail}") from exc


def discord_delete_message(token: str, channel_id: str, message_id: str) -> None:
    try:
        discord_request("DELETE", token, f"/channels/{channel_id}/messages/{message_id}")
    except RuntimeError as exc:
        print(f"Could not delete old media catalog message {message_id}: {exc}", file=sys.stderr)


def discord_bot_identity(token: str) -> dict[str, Any]:
    return discord_request("GET", token, "/users/@me")