nakama/moderation/backend-py/websocket_server.py

"""
WebSocket server for moderation chat
"""
import socketio
import logging
from typing import Dict, Set, Optional
from datetime import datetime

from config import settings
from database import moderation_admins_collection
from utils.auth import normalize_username

logger = logging.getLogger(__name__)

# Create Socket.IO server
sio = socketio.AsyncServer(
    async_mode='asgi',
    cors_allowed_origins='*',
    logger=False,
    engineio_logger=False
)

# Track connected moderators
connected_moderators: Dict[str, dict] = {}  # {sid: {username, telegramId, isOwner}}


def broadcast_online():
    """Broadcast list of online moderators"""
    unique = {}
    for data in connected_moderators.values():
        username = data.get('username')
        if username:
            unique[username] = data
    
    online_list = list(unique.values())
    sio.emit('online', online_list, namespace='/mod-chat')


# Namespace handlers for /mod-chat
@sio.on('connect', namespace='/mod-chat')
async def on_connect(sid, environ):
    """Handle client connection to /mod-chat namespace"""
    logger.info(f"[WebSocket] Client connected to /mod-chat: {sid}")
    # Don't authorize immediately - wait for 'auth' event


@sio.on('disconnect', namespace='/mod-chat')
async def on_disconnect(sid):
    """Handle client disconnection"""
    logger.info(f"[WebSocket] Client disconnected from /mod-chat: {sid}")
    
    if sid in connected_moderators:
        del connected_moderators[sid]
        broadcast_online()


@sio.on('auth', namespace='/mod-chat')
async def on_auth(sid, data):
    """Handle authentication for moderation chat"""
    try:
        username = normalize_username(data.get('username')) if data.get('username') else None
        telegram_id = data.get('telegramId')
        
        if not username or not telegram_id:
            logger.warning(f"[WebSocket] Auth failed: missing username or telegramId")
            await sio.emit('unauthorized', namespace='/mod-chat', room=sid)
            await sio.disconnect(sid, namespace='/mod-chat')
            return
        
        # Check if user is owner
        owner_usernames = settings.OWNER_USERNAMES_LIST
        is_owner = username.lower() in [u.lower() for u in owner_usernames]
        
        # Check if user is moderation admin
        admin = await moderation_admins_collection().find_one({
            'telegramId': str(telegram_id)
        })
        is_admin = admin is not None
        
        if not is_owner and not is_admin:
            logger.warning(f"[WebSocket] Access denied: {username} (telegramId: {telegram_id})")
            await sio.emit('unauthorized', namespace='/mod-chat', room=sid)
            await sio.disconnect(sid, namespace='/mod-chat')
            return
        
        # Store connection data
        connected_moderators[sid] = {
            'username': username,
            'telegramId': telegram_id,
            'isOwner': is_owner
        }
        
        logger.info(f"[WebSocket] Auth success: {username} (owner: {is_owner}, admin: {is_admin})")
        await sio.emit('ready', namespace='/mod-chat', room=sid)
        broadcast_online()
        
    except Exception as e:
        logger.error(f"[WebSocket] Error in auth: {e}")
        import traceback
        traceback.print_exc()
        await sio.emit('unauthorized', namespace='/mod-chat', room=sid)
        await sio.disconnect(sid, namespace='/mod-chat')


@sio.on('message', namespace='/mod-chat')
async def on_message(sid, data):
    """Handle moderation chat message"""
    try:
        if sid not in connected_moderators:
            logger.warning(f"[WebSocket] Message from unauthorized client: {sid}")
            return
        
        user_data = connected_moderators[sid]
        text = (data.get('text') or '').strip()
        
        if not text:
            return
        
        message = {
            'id': f"{int(datetime.utcnow().timestamp() * 1000)}-{sid[:8]}",
            'username': user_data['username'],
            'telegramId': user_data['telegramId'],
            'text': text,
            'createdAt': datetime.utcnow().isoformat()
        }
        
        # Broadcast to all in namespace
        await sio.emit('message', message, namespace='/mod-chat')
        logger.info(f"[WebSocket] Message from {user_data['username']}: {text[:50]}...")
        
    except Exception as e:
        logger.error(f"[WebSocket] Error handling message: {e}")
        import traceback
        traceback.print_exc()


def get_socketio_app():
    """Get Socket.IO ASGI app"""
    return socketio.ASGIApp(sio)