134 lines
4.1 KiB
Python
134 lines
4.1 KiB
Python
"""SocketIO service for real-time communication."""
|
|
|
|
import logging
|
|
|
|
from flask import request
|
|
from flask_jwt_extended import decode_token
|
|
from flask_socketio import disconnect, emit, join_room, leave_room
|
|
|
|
from app import socketio
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SocketIOService:
|
|
"""Service for managing SocketIO connections and user rooms."""
|
|
|
|
@staticmethod
|
|
def get_user_room(user_id: int) -> str:
|
|
"""Get the room name for a specific user."""
|
|
return f"user_{user_id}"
|
|
|
|
@staticmethod
|
|
def emit_to_user(user_id: int, event: str, data: dict) -> None:
|
|
"""Emit an event to a specific user's room."""
|
|
room = SocketIOService.get_user_room(user_id)
|
|
socketio.emit(event, data, room=room)
|
|
logger.debug(f"Emitted {event} to user {user_id} in room {room}")
|
|
|
|
@staticmethod
|
|
def emit_credits_changed(user_id: int, new_credits: int) -> None:
|
|
"""Emit credits_changed event to a user."""
|
|
SocketIOService.emit_to_user(
|
|
user_id,
|
|
"credits_changed",
|
|
{"credits": new_credits},
|
|
)
|
|
|
|
@staticmethod
|
|
def get_user_from_socketio() -> dict | None:
|
|
"""Get user from SocketIO connection using cookies."""
|
|
try:
|
|
from flask import current_app
|
|
from flask_jwt_extended import decode_token
|
|
|
|
# Check if we have the access_token cookie
|
|
access_token = request.cookies.get("access_token_cookie")
|
|
if not access_token:
|
|
return None
|
|
|
|
# Decode the JWT token manually
|
|
with current_app.app_context():
|
|
try:
|
|
decoded_token = decode_token(access_token)
|
|
current_user_id = decoded_token["sub"]
|
|
if not current_user_id:
|
|
return None
|
|
except Exception:
|
|
return None
|
|
|
|
# Query database for user data
|
|
from app.models.user import User
|
|
user = User.query.get(int(current_user_id))
|
|
if not user or not user.is_active:
|
|
return None
|
|
|
|
return {
|
|
"id": str(user.id),
|
|
"email": user.email,
|
|
"name": user.name,
|
|
"role": user.role,
|
|
"credits": user.credits,
|
|
}
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
@socketio.on("connect")
|
|
def handle_connect(auth=None) -> bool:
|
|
"""Handle client connection."""
|
|
try:
|
|
logger.info(f"SocketIO connection from {request.remote_addr}")
|
|
return True
|
|
|
|
except Exception:
|
|
logger.exception("Error handling SocketIO connection")
|
|
disconnect()
|
|
return False
|
|
|
|
|
|
@socketio.on("authenticate")
|
|
def handle_authenticate(data):
|
|
"""Handle authentication after connection."""
|
|
try:
|
|
user = SocketIOService.get_user_from_socketio()
|
|
if not user:
|
|
emit("auth_error", {"error": "Authentication failed"})
|
|
disconnect()
|
|
return
|
|
|
|
user_id = int(user["id"])
|
|
user_room = SocketIOService.get_user_room(user_id)
|
|
|
|
# Join user-specific room
|
|
join_room(user_room)
|
|
|
|
logger.info(f"User {user_id} authenticated and joined room {user_room}")
|
|
|
|
# Send current credits on authentication
|
|
emit("auth_success", {"user": user})
|
|
emit("credits_changed", {"credits": user["credits"]})
|
|
|
|
except Exception:
|
|
logger.exception("Error handling SocketIO authentication")
|
|
emit("auth_error", {"error": "Authentication failed"})
|
|
disconnect()
|
|
|
|
|
|
@socketio.on("disconnect")
|
|
def handle_disconnect() -> None:
|
|
"""Handle client disconnection."""
|
|
try:
|
|
user = SocketIOService.get_user_from_socketio()
|
|
if user:
|
|
user_id = int(user["id"])
|
|
user_room = SocketIOService.get_user_room(user_id)
|
|
leave_room(user_room)
|
|
logger.info(f"User {user_id} disconnected from SocketIO")
|
|
|
|
except Exception:
|
|
logger.exception("Error handling SocketIO disconnection")
|
|
|
|
|
|
# Export the service instance
|
|
socketio_service = SocketIOService() |