308 lines
13 KiB
Python
308 lines
13 KiB
Python
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()
|