Connect to Discord Gateway via websockets. Statically typed, slightly optimized, and color coded.

Discord-Gateway.py

import asyncio
import json
import logging
from typing import NoReturn

import websockets
import websockets.exceptions
from colorama import Fore, init

init(autoreset=True)

discord_ws_url = "wss://gateway.discord.gg/?v=10&encoding=json"
TOKEN = "INSERT-YOUR-TOKEN"

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[logging.StreamHandler()],
)


async def listener() -> NoReturn:
    """Listen to Discord's WebSocket gateway for events."""
    while True:
        try:
            async with websockets.connect(
                discord_ws_url, ping_interval=20, ping_timeout=20
            ) as ws:
                await identify(ws)
                await on_message(ws)
        except websockets.exceptions.ConnectionClosed as e:
            logging.error(
                f"{Fore.RED}WebSocket connection closed unexpectedly: {e}. Reconnecting..."
            )
            await asyncio.sleep(5)
        except Exception as err:
            logging.error(f"{Fore.RED}An error occurred: {err}. Reconnecting...")
            await asyncio.sleep(5)


async def identify(ws: websockets.WebSocketClientProtocol) -> NoReturn:
    """Identify the bot with Discord gateway."""
    logging.info(f"{Fore.GREEN}Identifying...")

    identify_payload: dict[str, int | dict[str, str | dict[str, str]]] = {
        "op": 2,
        "d": {
            "token": TOKEN,
            "properties": {
                "$os": "windows",
                "$browser": "chrome",
                "$device": "pc",
            },
        },
    }

    await ws.send(json.dumps(identify_payload))
    logging.info(f"{Fore.CYAN}Identification sent.")


async def heartbeat(
    ws: websockets.WebSocketClientProtocol, interval: float, last_sequence: int | None
) -> NoReturn:
    """Send heartbeat to keep the connection alive."""
    while True:
        await asyncio.sleep(interval)
        heartbeat_payload: dict[str, int | None] = {"op": 1, "d": last_sequence}
        await ws.send(json.dumps(heartbeat_payload))
        logging.debug(f"{Fore.YELLOW}Heartbeat sent.")


def log_event(event_type: str | None, event: dict) -> None:
    """Log events with color-coded output."""
    color_map: dict[str | None, str] = {
        "MESSAGE_CREATE": Fore.LIGHTBLUE_EX,
        "MESSAGE_UPDATE": Fore.LIGHTCYAN_EX,
        "PRESENCE_UPDATE": Fore.LIGHTMAGENTA_EX,
        "READY": Fore.LIGHTGREEN_EX,
        None: Fore.LIGHTYELLOW_EX,
    }
    color: str = color_map.get(event_type, Fore.LIGHTWHITE_EX)
    logging.info(f"{color}Event received: {event['t']} {event.get('d', {})}")


async def on_message(ws: websockets.WebSocketClientProtocol) -> NoReturn:
    """Handle incoming messages from the WebSocket."""
    logging.info(f"{Fore.GREEN}Listening for messages...")
    last_sequence = None
    while True:
        event = await ws.recv()
        event = json.loads(event)

        op_code = event.get("op")
        event_type = event.get("t")

        log_event(event_type, event)

        if op_code == 10:
            # Handle opcode 10 (Hello)
            interval = event["d"]["heartbeat_interval"] / 1000
            asyncio.create_task(heartbeat(ws, interval, last_sequence))

        elif op_code == 0:
            # Handle opcode 0 (Dispatch)
            last_sequence = event.get("s")

            if event_type == "MESSAGE_CREATE":
                # Handle the MESSAGE_CREATE event
                # You can get metadata of a message in event["d"]
                logging.info(f"{Fore.LIGHTBLUE_EX}{event['d'].get('content', '')}")

        elif op_code == 9:
            # Handle opcode 9 (Invalid Session)
            logging.info(f"{Fore.RED}Invalid session. Starting a new session...")
            await identify(ws)


async def main() -> NoReturn:
    """Main entry point for the bot."""
    await listener()


if __name__ == "__main__":
    asyncio.run(main())