k.eaven
5669048f89
- Core infrastructure: config, state management, models - REST API server (aiohttp) with Discord-compatible endpoints - WebSocket Gateway server (websockets) with Discord Gateway protocol - State management with guilds, channels, users, messages, interactions - Pydantic models for Discord objects (User, Guild, Channel, Message, etc.) - CLI with start, test, reset, export, validate commands - YAML-based test scenario framework - Basic test scenario for /server_status command flow - Unit tests for models and state manager (19 passing) - GPL-3.0 license
285 lines
9.0 KiB
Python
285 lines
9.0 KiB
Python
"""
|
|
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") |