Initial commit: APIcord Discord API Emulator
- 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
This commit is contained in:
k.eaven committed
commit
5669048f89
38 files changed
+2871
No files matched your search
@@ -0,0 +1,285 @@
|
||||
"""
|
||||
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")
|
||||
Reference in new issue
Block a user