""" WebSocket Gateway Server Module Implements the Discord Gateway WebSocket server for real-time events. """ import asyncio import json import logging import random import time from typing import Any, Dict, Optional, Set import websockets from websockets.server import WebSocketServerProtocol from apicord.config import EmulatorConfig from apicord.state.manager import StateManager from apicord.models import ( User, Guild, Channel, Role, Interaction, InteractionType, GatewayInfo, ) logger = logging.getLogger("apicord.gateway") class GatewayConnection: """Represents a single WebSocket connection to the gateway.""" def __init__( self, websocket: WebSocketServerProtocol, config: EmulatorConfig, state: StateManager, ) -> None: self.ws = websocket self.config = config self.state = state self.sequence: Optional[int] = None self.session_id: Optional[str] = None self.heartbeat_interval: float = 41250 / 1000 # 41.25 seconds (Discord default) self.heartbeat_task: Optional[asyncio.Task] = None self.identified: bool = False self.user: Optional[User] = None self.guilds: Set[str] = set() self._last_heartbeat_ack: bool = True async def handle(self) -> None: """Main connection handler loop.""" logger.info(f"New gateway connection from {self.ws.remote_address}") try: # Send HELLO await self._send_hello() # Process incoming messages async for message in self.ws: try: data = json.loads(message) await self._handle_message(data) except json.JSONDecodeError: logger.warning(f"Invalid JSON from {self.ws.remote_address}") except Exception as e: logger.error(f"Error handling message: {e}", exc_info=True) except websockets.exceptions.ConnectionClosed: logger.info(f"Connection closed: {self.ws.remote_address}") except Exception as e: logger.error(f"Connection error: {e}", exc_info=True) finally: await self._cleanup() async def _send_hello(self) -> None: """Send HELLO payload (opcode 10).""" hello = { "op": 10, "d": { "heartbeat_interval": int(self.heartbeat_interval * 1000), } } await self._send(hello) logger.debug("Sent HELLO") # Start heartbeat monitoring self.heartbeat_task = asyncio.create_task(self._heartbeat_monitor()) async def _heartbeat_monitor(self) -> None: """Monitor heartbeats and send heartbeats.""" while True: try: await asyncio.sleep(self.heartbeat_interval) if not self._last_heartbeat_ack: logger.warning("Heartbeat ACK not received, closing connection") await self.ws.close(code=4009, reason="Heartbeat timeout") break self._last_heartbeat_ack = False await self._send({"op": 1, "d": self.sequence}) except asyncio.CancelledError: break except Exception as e: logger.error(f"Heartbeat error: {e}") break async def _handle_message(self, data: Dict[str, Any]) -> None: """Handle incoming gateway message.""" op = data.get("op") payload = data.get("d") seq = data.get("s") event_name = data.get("t") # Update sequence number if seq is not None: self.sequence = seq if op == 1: # Heartbeat await self._handle_heartbeat() elif op == 2: # Identify await self._handle_identify(payload) elif op == 3: # Presence Update await self._handle_presence_update(payload) elif op == 4: # Voice State Update pass elif op == 6: # Resume await self._handle_resume(payload) elif op == 11: # Heartbeat ACK self._last_heartbeat_ack = True else: logger.warning(f"Unknown opcode: {op}") async def _handle_heartbeat(self) -> None: """Handle incoming heartbeat (opcode 1).""" self._last_heartbeat_ack = True await self._send({"op": 11}) # Heartbeat ACK logger.debug("Sent heartbeat ACK") async def _handle_identify(self, payload: Dict[str, Any]) -> None: """Handle IDENTIFY payload (opcode 2).""" token = payload.get("token") if token != self.config.bot_token: logger.warning("Invalid token in IDENTIFY") await self.ws.close(code=4004, reason="Invalid token") return # Mark as identified self.identified = True self.session_id = f"session_{random.randint(100000, 999999)}" # Get bot user self.user = self.state.get_bot_user() # Send READY event ready_data = { "v": self.config.gateway_version, "user": self.user.model_dump() if self.user else {}, "guilds": [ {"id": g.id, "unavailable": False} for g in self.state.list_guilds() ], "session_id": self.session_id, "shard": [0, 1], "application": { "id": self.config.bot_id, "flags": 0 } } await self._send({ "op": 0, "t": "READY", "s": self.sequence, "d": ready_data }) logger.info(f"Bot identified: {self.user.username if self.user else 'Unknown'} ({self.session_id})") async def _handle_presence_update(self, payload: Dict[str, Any]) -> None: """Handle presence update.""" # Just acknowledge for now pass async def _handle_resume(self, payload: Dict[str, Any]) -> None: """Handle RESUME payload (opcode 6).""" token = payload.get("token") session_id = payload.get("session_id") seq = payload.get("seq") if token != self.config.bot_token or session_id != self.session_id: logger.warning("Invalid RESUME") await self.ws.close(code=4004, reason="Invalid session") return self.identified = True self.sequence = seq # Send RESUMED event await self._send({ "op": 0, "t": "RESUMED", "s": self.sequence, "d": {} }) logger.info("Session resumed") async def _send(self, payload: Dict[str, Any]) -> None: """Send JSON payload over WebSocket.""" await self.ws.send(json.dumps(payload)) async def send_event(self, event_name: str, data: Dict[str, Any]) -> None: """Send a gateway event (opcode 0).""" if not self.identified: return self.sequence = (self.sequence or 0) + 1 payload = { "op": 0, "t": event_name, "s": self.sequence, "d": data } await self._send(payload) async def _cleanup(self) -> None: """Cleanup on connection close.""" if self.heartbeat_task: self.heartbeat_task.cancel() try: await self.heartbeat_task except asyncio.CancelledError: pass class GatewayServer: """Discord Gateway WebSocket server.""" def __init__(self, config: EmulatorConfig, state: StateManager) -> None: self.config = config self.state = state self.server: Optional[websockets.WebSocketServer] = None self.connections: Set[GatewayConnection] = set() async def start(self) -> None: """Start the WebSocket gateway server.""" self.server = await websockets.serve( self._handle_connection, self.config.host, self.config.gateway_port, ping_interval=20, ping_timeout=10, ) logger.info(f"Gateway server started on ws://{self.config.host}:{self.config.gateway_port}") async def _handle_connection(self, websocket: WebSocketServerProtocol) -> None: """Handle new WebSocket connection.""" connection = GatewayConnection(websocket, self.config, self.state) self.connections.add(connection) try: await connection.handle() finally: self.connections.discard(connection) def broadcast_event(self, event_name: str, data: Dict[str, Any]) -> None: """Broadcast an event to all connected clients.""" for conn in self.connections: if conn.identified: asyncio.create_task(conn.send_event(event_name, data)) async def stop(self) -> None: """Stop the gateway server.""" # Close all connections for conn in self.connections: await conn.ws.close(code=1001, reason="Server shutting down") if self.server: self.server.close() await self.server.wait_closed() logger.info("Gateway server stopped")