import json from typing import Dict, List, Union from fastapi import Cookie, Depends, FastAPI, HTTPException, Query, WebSocket, status, APIRouter, WebSocketDisconnect, status from fastapi.responses import HTMLResponse from config import User_schema from database.auth.models import UserModel from database.exercices.crud import generate_unique_code from database.room.crud import check_anonymous_owner, check_user_in_room, check_user_owner, connect_room, create_waiter_anonymous, create_waiter_by_user, disconnect_room, get_member_by_code, validate_name_in_room from database.room.models import AnonymousMember, Room, RoomOwner, Waiter from services.auth import get_user_from_token from services.io import get_abs_path_from_relative_to_root import secrets router = APIRouter() @router.get('/') def index(): return HTMLResponse(open(get_abs_path_from_relative_to_root('/index.html'), 'r').read()) class ConnectionManager: def __init__(self): self.active_connections: Dict[str,List[WebSocket]] = {} async def add(self, group, ws): if group not in self.active_connections: self.active_connections[group] = [] if ws not in self.active_connections[group]: self.active_connections[group].append(ws) def remove(self, ws: WebSocket, group): if group in self.active_connections: if ws in self.active_connections[group]: self.active_connections[group].remove(ws) async def send_personal_message(self, message: str, websocket: WebSocket): await websocket.send_text(message) async def broadcast(self, message: str, group): if group in self.active_connections: for connection in self.active_connections[group]: await connection.send_json(message) manager = ConnectionManager() class Consumer(): def __init__(self, ws: WebSocket): self.ws : WebSocket = ws async def connect(self): pass async def receive(self): pass async def disconnect(self): pass async def run(self): await self.connect() try: while True: data = await self.ws.receive_text() await self.receive(data) except WebSocketDisconnect: await self.disconnect() class RoomConsumer(Consumer): def __init__(self, ws:WebSocket, room_id, manager:ConnectionManager): self.ws:WebSocket = ws self.room_id = room_id self.manager:ConnectionManager = manager self.owner = False async def connect(self): await self.ws.accept() self.clientId = secrets.token_hex(32) await self.manager.add(self.ws, self.room_id) self.room : Room = await Room.get(id_code=self.room_id) self.status = None self.waiter = None self.user = None await self.ws.send_json({'type': 'accept'}) async def receive(self, data): json_data = json.loads(data) payload = json_data['data'] type = json_data['type'] if type == 'auth': token = payload['token'] self.user = await get_user_from_token(token) if self.user is not None: await self.ws.send_json({'type': 'auth_success'}) else: await self.ws.send_json({'type': 'auth_failed'}) if type == "login" and self.room.private == True and self.user is not None: if await check_user_in_room(self.room, self.user): if await check_user_owner(self.room, self.user): self.owner = True await self.manager.add(f'{self.room_id}__owner', self.ws) else: await self.manager.add(self.room_id, self.ws) await connect_room(self.room, self.user.id) await self.manager.broadcast({'type': 'joined', 'data': {"name": self.user.username}}, self.room_id) await self.manager.broadcast({'type': 'joined', 'data': {"name": self.user.username}}, f'{self.room_id}__owner') else: self.waiter = await create_waiter_by_user(self.room, self.user) await self.ws.send_json({'data': "waiting"}) await self.manager.add(f'{self.room_id}__waiting__{self.user.id}', self.ws) await self.manager.broadcast({'type': 'add_waiter', 'data': {"name": self.user.username, 'id': self.user.id}}, f'{self.room_id}__owner') if type == "login" and self.room.private == True and self.user is None: if 'relogin_code' in payload: anonymous = await get_member_by_code(self.room, payload['relogin_code']) if anonymous is not None: if await check_anonymous_owner(self.room, anonymous): self.owner = True await self.manager.add(f'{self.room_id}__owner', self.ws) else: await self.manager.add(self.room_id, self.ws) self.anonymous = anonymous await connect_room(self.room, self.anonymous.id_code) await self.manager.broadcast( {'type': 'joined', 'data': {"name": anonymous.name}}, self.room_id) else: valid_username = await validate_name_in_room(payload['name']) if valid_username == True: self.waiter = await create_waiter_anonymous(self.room, payload['name']) await self.ws.send_json({'type': "waiting"}) await self.manager.add(f'{self.room_id}__waiting__{self.waiter.id_code}', self.ws) await self.manager.broadcast({'type': 'add_waiter', 'data': { "name": self.waiter.name, 'id': self.waiter.id_code}}, f'{self.room_id}__owner') else: await self.ws.send_json({"type": "error", 'data': {"user_input": valid_username}}) if type == "login" and self.room.private == False and self.user is not None: if await check_user_in_room(self.room, self.user): if await check_user_owner(self.room, self.user): self.owner = True await self.manager.add(f'{self.room_id}__owner', self.ws) else: await self.manager.add(self.room_id, self.ws) await connect_room(self.room, self.user.id) await self.manager.broadcast({'type': 'joined', 'data': {"name": self.user.username}}, self.room_id) await self.manager.broadcast({'type': 'joined', 'data': {"name": self.user.username}}, f'{self.room_id}__owner') else: await self.room.users.add(self.user) await self.manager.add(self.room_id, self.ws) await connect_room(self.room, self.user.id) await self.manager.broadcast( {'type': 'joined', 'data': {"name": self.user.username}}, self.room_id) await self.manager.broadcast( {'type': 'joined', 'data': {"name": self.user.username}}, f'{self.room_id}__owner') if type == 'login' and self.room.private == False and self.user is None: if 'relogin_code' in payload: anonymous = await get_member_by_code(self.room, payload['relogin_code']) if anonymous is not None: if await check_anonymous_owner(self.room, anonymous): self.owner = True await self.manager.add(f'{self.room_id}__owner', self.ws) else: await self.manager.add(self.room_id, self.ws) self.anonymous = anonymous await connect_room(self.room, self.anonymous.id_code) await self.manager.broadcast( {'type': 'joined', 'data': {"name": anonymous.name}}, self.room_id) else: valid_username = await validate_name_in_room(self.room, payload['name']) if valid_username == True: code = await generate_unique_code(AnonymousMember) self.anonymous = await AnonymousMember.create(name=payload['name'], id_code=code, room_id=self.room.id) await self.manager.add(self.room_id, self.ws) self.owner = False await connect_room(self.room, self.anonymous.id_code) await self.manager.broadcast({'type': 'joined', 'data': {'name': self.anonymous.name}}, self.room_id) await self.manager.broadcast({'type': 'joined', 'data': {'name': self.anonymous.name, 'code': self.anonymous.id_code}}, f'{self.room_id}__owner') else: await self.ws.send({'type': "error", "data": {"user_input": valid_username}}) if type == 'accept_waiter': if self.owner == True: id = payload['id'] await self.manager.broadcast({'type': 'log_waiter', 'data': {}}, f'{self.room_id}__waiting__{id}') if type == 'refuse_waiter': if self.owner == True: id = payload['id'] await self.manager.broadcast({'type': 'reject_waiter', 'data': {}}, f'{self.room_id}__waiting__{id}') if type == 'log_waiter': if self.user is not None: await self.room.users.add(self.user) await self.waiter.delete() self.manager.remove('f{self.room_id}__waiting__{self.user.id}', self.ws) await self.manager.add(self.room_id, self.ws) await connect_room(self.room, self.user.id) await self.manager.broadcast( {'type': 'joined', 'data': {"name": self.user.username}}, self.room_id) await self.manager.broadcast( {'type': 'joined', 'data': {"name": self.user.username}}, f'{self.room_id}__owner') else: code = await generate_unique_code(AnonymousMember) self.anonymous = await AnonymousMember.create(name=self.waiter.name, id_code=code, room_id=self.room.id) self.manager.remove(self.ws, f'{self.room_id}__waiting__{self.waiter.id_code}') await self.waiter.delete() self.waiter = None await self.manager.add(self.room_id, self.ws) self.owner = False await connect_room(self.room, self.anonymous.id_code) await self.manager.broadcast({'type': 'joined', 'data': {'name': self.anonymous.name}}, self.room_id) await self.manager.broadcast({'type': 'joined', 'data': {'name': self.anonymous.name, 'code': self.anonymous.id_code}}, f'{self.room_id}__owner') elif type == 'ban': if self.owner == True: status = payload['status'] name = "" if status == 'user': user = await UserModel.get(id=payload['id']) await disconnect_room(self.room, user.id) name = user.username await self.room.users.remove(user) elif status == 'anonymous': anonymous = await AnonymousMember.get(id_code=payload['id']) name = anonymous.name await disconnect_room(self.room, anonymous.id_code) await anonymous.delete() await self.manager.broadcast({'type': 'leave', 'data': {"name": name}}, self.room_id) await self.manager.broadcast({'type': 'leave', 'data': {"name": name}},f'{self.room_id}__owner') elif type == "leave": name = "" if self.user is not None: name = self.user.username await self.room.users.remove(self.user) else: name = self.anonymous.name await self.anonymous.delete() await self.manager.broadcast({'type': 'leave', 'data': {"name": name}}, self.room_id) await self.manager.broadcast({'type': 'leave', 'data': {"name": name}}, f'{self.room_id}__owner') async def disconnect(self): if self.waiter != None: self.manager.remove(self.ws, f'{self.room_id}__waiting__{self.waiter.id_code}') await self.manager.broadcast({'type': "disconnect_waiter", 'data': {'name': self.waiter.name}}, self.room_id) await self.manager.broadcast({'type': "disconnect_waiter", 'data': {'name': self.waiter.name}}, f'{self.room_id}__owner') await self.waiter.delete() if self.owner == True: if self.user is not None: disconnect_room(self.room, self.user.id) else: disconnect_room(self.room, self.anonymous.id_code) self.manager.remove(self.ws, f'{self.room_id}__owner') else: self.manager.remove(self.ws, self.room_id) if self.user is not None: disconnect_room(self.room, self.user.id) await self.manager.broadcast({'type': "disconnect", 'data': {'name': self.user.username}}, self.room_id) await self.manager.broadcast({'type': "disconnect", 'data': {'name': self.user.username}}, f'{self.room_id}__owner') elif self.anonymous is not None: disconnect_room(self.room, self.anonymous.id_code) await self.manager.broadcast({'type': "disconnect_waiter", 'data': {'name': self.anonymous.name}}, self.room_id) await self.manager.broadcast({'type': "disconnect_waiter", 'data': {'name': self.anonymous.name}}, f'{self.room_id}__owner') await self.manager.broadcast({'type': "disconnect", 'data': {'name': self}}, self.room_id) await self.manager.broadcast({'type': f"Client left the chat"}, f'{self.room_id}__owner') async def check_room(room_id): room = await Room.get_or_none(id_code = room_id) if room == None: raise HTTPException(status_code = status.HTTP_404_NOT_FOUND, detail = 'Room does not exist ') return room_id @router.websocket('/ws/{room_id}') async def room_ws(ws: WebSocket, room_id:str = Depends(check_room), ): consumer = RoomConsumer(ws, room_id, manager) await consumer.run()