556 lines
22 KiB
Python
556 lines
22 KiB
Python
from database.room.crud import create_room_db, delete_member, get_member, get_member_from_token, get_member_from_reconnect_code, connect_member, disconnect_member, create_anonymous_member, create_anonymous_waiter, create_user_member, create_user_waiter, get_waiter, accept_waiter, leave_room, refuse_waiter, check_room
|
|
from pydantic.error_wrappers import ValidationError
|
|
from pydantic import validate_arguments
|
|
from sqlmodel import select
|
|
from pydantic import BaseModel
|
|
from typing import Any, Callable, Dict, List, Optional
|
|
from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect, status, Query
|
|
from config import ALGORITHM, SECRET_KEY
|
|
from database.auth.crud import get_user_from_clientId_db
|
|
from database.auth.models import User
|
|
from database.db import get_session
|
|
|
|
from sqlmodel import Session
|
|
from database.room.models import Anonymous, MemberSerializer, MemberWithRelations, MemberRead, Room, RoomCreate, RoomRead, Member
|
|
from services.auth import get_current_user_optional
|
|
from fastapi.exceptions import HTTPException
|
|
from jose import jwt, exceptions
|
|
import inspect
|
|
router = APIRouter(tags=["room"])
|
|
|
|
|
|
class RoomAndMember(BaseModel):
|
|
room: RoomRead
|
|
member: MemberRead
|
|
|
|
|
|
class Waiter(BaseModel):
|
|
username: str
|
|
waiter_id: str
|
|
|
|
|
|
def serialize_member(member: Member) -> MemberRead | Waiter:
|
|
member_obj = member.user or member.anonymous
|
|
if member.waiting == False:
|
|
return MemberRead(username=member_obj.username, reconnect_code=getattr(member_obj, "reconnect_code", ""), isUser=member.user_id != None, isAdmin=member.is_admin, id_code=member.id_code)
|
|
if member.waiting == True:
|
|
return Waiter(username=member_obj.username, waiter_id=member.id_code)
|
|
|
|
|
|
@router.post('/room', response_model=RoomAndMember)
|
|
def create_room(room: RoomCreate, username: Optional[str] = Query(default=None, max_length=20), user: User | None = Depends(get_current_user_optional), db: Session = Depends(get_session)):
|
|
room_obj = create_room_db(room=room, user=user, username=username, db=db)
|
|
return {'room': room_obj['room'], "member": serialize_member(room_obj['member'])}
|
|
|
|
|
|
class ConnectionManager:
|
|
def __init__(self):
|
|
self.active_connections: Dict[str, List[WebSocket]] = {}
|
|
|
|
def add(self, group: str, ws: WebSocket):
|
|
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, group: str, ws: WebSocket):
|
|
if group in self.active_connections:
|
|
if ws in self.active_connections[group]:
|
|
self.active_connections[group].remove(ws)
|
|
|
|
async def broadcast(self, message, group: str, exclude: list[WebSocket] = []):
|
|
if group in self.active_connections:
|
|
for connection in self.active_connections[group]:
|
|
if connection not in exclude:
|
|
await connection.send_json(message)
|
|
|
|
|
|
def make_event_decorator(eventsDict):
|
|
def _(name: str | List, conditions: List[Callable | bool] = []):
|
|
def add_event(func):
|
|
model = validate_arguments(func).model
|
|
if type(name) == str:
|
|
eventsDict[name] = {"func": func,
|
|
"conditions": conditions, "model": model}
|
|
if type(name) == list:
|
|
for n in name:
|
|
eventsDict[n] = {"func": func,
|
|
"conditions": conditions, "model": model}
|
|
return func
|
|
return add_event
|
|
return _
|
|
|
|
|
|
class Event(BaseModel):
|
|
func: Callable
|
|
conditions: List[Callable | bool]
|
|
model: BaseModel
|
|
|
|
|
|
def dict_model(model: BaseModel, exclude: List[str]):
|
|
value = {}
|
|
for n, f in model:
|
|
if n not in exclude:
|
|
value[n] = f
|
|
return value
|
|
|
|
|
|
def dict_all(obj: Any):
|
|
if isinstance(obj, dict):
|
|
value = {}
|
|
for k, v in obj.items():
|
|
if isinstance(v, dict):
|
|
v = dict_all(v)
|
|
value[k] = dict(v)
|
|
elif isinstance(v, BaseModel):
|
|
value[k] = dict(v)
|
|
else:
|
|
try:
|
|
value[k] = dict(v)
|
|
except:
|
|
value[k] = v
|
|
return value
|
|
return dict(obj)
|
|
|
|
|
|
class Consumer:
|
|
events: Dict[str, Event] = {}
|
|
sendings: Dict[str, Any] = {}
|
|
event = make_event_decorator(events)
|
|
sending = make_event_decorator(sendings)
|
|
|
|
def __init__(self, ws: WebSocket):
|
|
self.ws: WebSocket = ws
|
|
#self.events: Dict[str, Callable] = {}
|
|
|
|
async def connect(self):
|
|
pass
|
|
|
|
async def validation_error_handler(self, e: ValidationError):
|
|
errors = e.errors()
|
|
await self.ws.send_json({"type": "error", "data": {"detail": [{ers['loc'][-1]: ers['msg']} for ers in errors]}})
|
|
|
|
async def send(self, payload):
|
|
type = payload.get('type', None)
|
|
print('TYPE', type, self.member)
|
|
if type is not None:
|
|
event_wrapper = self.sendings.get(type, None)
|
|
if event_wrapper is not None:
|
|
handler = event_wrapper.get('func')
|
|
conditions = event_wrapper.get('conditions')
|
|
|
|
is_valid = all([(await c(self)) if inspect.iscoroutinefunction(c) else c(self) if inspect.isfunction(c) else c == True if isinstance(c, bool) else True for c in conditions])
|
|
|
|
if handler is not None and is_valid:
|
|
model = event_wrapper.get("model")
|
|
|
|
data = payload.get('data') or {}
|
|
try:
|
|
validated_payload = model(self=self, **data)
|
|
except ValidationError as e:
|
|
print("ERROR", e)
|
|
await self.ws.send_json({"type": "error", "data": {"msg": "Oops there was an error"}})
|
|
return
|
|
|
|
validated_payload = dict_model(validated_payload,
|
|
exclude=["v__duplicate_kwargs", "args", 'kwargs', "self"])
|
|
try:
|
|
parsed_payload = handler(
|
|
self, **validated_payload)
|
|
|
|
await self.ws.send_json({'type': type, "data": dict_all(parsed_payload)})
|
|
return
|
|
except Exception as e:
|
|
return
|
|
return
|
|
#print('pls')
|
|
await self.ws.send_json(payload)
|
|
#print('sent')
|
|
|
|
async def receive(self, data):
|
|
event = data.get('type', None)
|
|
if event is not None:
|
|
event_wrapper = self.events.get(event, None)
|
|
if event_wrapper is not None:
|
|
handler = event_wrapper.get('func')
|
|
conditions = event_wrapper.get('conditions')
|
|
|
|
is_valid = all([(await c(self)) if inspect.iscoroutinefunction(c) else c(self) if inspect.isfunction(c) else c == True if isinstance(c, bool) else True for c in conditions])
|
|
|
|
if handler is not None and is_valid:
|
|
model = event_wrapper.get("model")
|
|
|
|
payload = data.get('data') or {}
|
|
try:
|
|
validated_payload = model(self=self, **payload)
|
|
except ValidationError as e:
|
|
await self.validation_error_handler(e)
|
|
return
|
|
|
|
await handler(**{k: v for k, v in validated_payload.dict().items() if k not in ["v__duplicate_kwargs", "args", 'kwargs']})
|
|
|
|
async def disconnect(self):
|
|
pass
|
|
|
|
async def run(self):
|
|
await self.connect()
|
|
try:
|
|
while True:
|
|
data = await self.ws.receive_json()
|
|
await self.receive(data)
|
|
except WebSocketDisconnect:
|
|
await self.disconnect()
|
|
|
|
|
|
class ConsumerManager:
|
|
def __init__(self):
|
|
self.active_connections: Dict[str, List[Consumer]] = {}
|
|
|
|
def add(self, group: str, ws: Consumer):
|
|
|
|
if group not in self.active_connections:
|
|
self.active_connections[group] = []
|
|
#print("adding", ws, self.active_connections[group])
|
|
if ws not in self.active_connections[group]:
|
|
#print('ACTUALLY ADDING')
|
|
self.active_connections[group].append(ws)
|
|
|
|
def remove(self, group: str, ws: Consumer):
|
|
if group in self.active_connections:
|
|
if ws in self.active_connections[group]:
|
|
self.active_connections[group].remove(ws)
|
|
|
|
async def broadcast(self, message, group: str, exclude: list[Consumer] = []):
|
|
if group in self.active_connections:
|
|
#print(self.active_connections[group], exclude)
|
|
for connection in list(set(self.active_connections[group])):
|
|
if connection not in exclude:
|
|
#print('SEND TO', connection, message)
|
|
await connection.send(message)
|
|
|
|
|
|
|
|
|
|
class Token(BaseModel):
|
|
token: str
|
|
|
|
|
|
def get_user_from_token(token: str, db: Session):
|
|
try:
|
|
decoded = jwt.decode(token=token, key=SECRET_KEY,
|
|
algorithms=[ALGORITHM])
|
|
except exceptions.ExpiredSignatureError:
|
|
return False
|
|
|
|
clientId = decoded.get('sub')
|
|
return get_user_from_clientId_db(clientId=clientId, db=db)
|
|
|
|
|
|
def check_same_member(member: Member, memberR: MemberRead):
|
|
return (member.user is not None and memberR.username == member.user.username) or (member.anonymous and memberR.reconnect_code == member.anonymous.reconnect_code)
|
|
|
|
class RoomConsumer(Consumer):
|
|
|
|
def __init__(self, ws: WebSocket, room: Room, manager: "RoomManager", db: Session):
|
|
self.room = room
|
|
self.ws = ws
|
|
self.manager = manager
|
|
self.db = db
|
|
self.member = None
|
|
# WS Utilities
|
|
|
|
async def connect(self):
|
|
await self.ws.accept()
|
|
|
|
async def direct_send(self, type: str, payload: Any):
|
|
await self.ws.send_json({'type': type, "data": payload})
|
|
|
|
async def send_to_admin(self, type: str, payload: Any, exclude: bool = False):
|
|
await self.manager.send_to_admin(self.room.id, {'type': type, "data": payload})
|
|
|
|
async def send_to(self, type: str, payload: Any,member_id, exclude: bool = False):
|
|
await self.manager.send_to(self.room.id, member_id,{'type': type, "data": payload})
|
|
|
|
async def broadcast(self, type, payload, exclude=False):
|
|
await self.manager.broadcast({"type": type, "data": payload}, self.room.id, exclude=[exclude == True and self])
|
|
|
|
def add_to_group(self):
|
|
self.manager.add(self.room.id, self)
|
|
|
|
async def connect_self(self):
|
|
if isinstance(self.member, Member):
|
|
connect_member(self.member, self.db)
|
|
await self.broadcast(type="connect", payload={"member": serialize_member(self.member).dict()}, exclude=True)
|
|
|
|
async def disconnect_self(self):
|
|
if isinstance(self.member, Member):
|
|
disconnect_member(self.member, self.db)
|
|
if self.member.waiting is False:
|
|
await self.broadcast(type="disconnect", payload={"member": serialize_member(self.member).dict()})
|
|
else:
|
|
await self.broadcast(type="disconnect_waiter", payload={"waiter": serialize_member(self.member).dict()})
|
|
|
|
# DB Utilities
|
|
|
|
# Received Events
|
|
@Consumer.event('login')
|
|
async def login(self, token: str | None = None, reconnect_code: str | None = None):
|
|
if token is not None:
|
|
member = get_member_from_token(token, self.room.id, self.db)
|
|
if member == False:
|
|
await self.direct_send(type="error", payload={"msg": "Token expired"})
|
|
return
|
|
if member is None:
|
|
await self.direct_send(type="error", payload={"msg": "Utilisateur introuvable dans cette salle"})
|
|
return
|
|
self.member = member
|
|
|
|
await self.connect_self()
|
|
self.add_to_group()
|
|
await self.direct_send(type="loggedIn", payload={"member": serialize_member(self.member).dict()})
|
|
|
|
elif reconnect_code is not None:
|
|
member = get_member_from_reconnect_code(
|
|
reconnect_code, self.room.id, db=self.db)
|
|
if member is None:
|
|
await self.direct_send(type="error", payload={"msg": "Utilisateur introuvable dans cette salle"})
|
|
return
|
|
|
|
self.member = member
|
|
|
|
await self.connect_self()
|
|
self.add_to_group()
|
|
await self.direct_send(type="loggedIn", payload={"member": serialize_member(self.member).dict()})
|
|
if reconnect_code is None and token is None:
|
|
await self.direct_send(type="error", payload={"msg": "Veuillez spécifier une méthode de connection"})
|
|
|
|
@Consumer.event('join')
|
|
async def join(self, token: str | None = None, username: str | None = None):
|
|
if self.room.public == False:
|
|
if token is not None:
|
|
user = get_user_from_token(token, self.db)
|
|
|
|
if user is None:
|
|
await self.direct_send(type="error", payload={"msg": "Utilisateur introuvable"})
|
|
return
|
|
if user is False:
|
|
await self.direct_send(type="error", payload={"msg": "Token expired"})
|
|
return
|
|
|
|
waiter = create_user_waiter(user, self.room, self.db)
|
|
|
|
if waiter.waiting is False:
|
|
self.member = waiter
|
|
# await self.connect_self()
|
|
self.add_to_group()
|
|
await self.connect_self()
|
|
await self.direct_send(type="loggedIn", payload={"member": serialize_member(self.member).dict()})
|
|
return
|
|
|
|
self.member = waiter
|
|
self.add_to_group()
|
|
await self.direct_send(type="waiting", payload={"waiter": serialize_member(self.member).dict()})
|
|
await self.send_to_admin(type="waiter", payload={"waiter": serialize_member(self.member).dict()})
|
|
|
|
elif username is not None:
|
|
waiter = create_anonymous_waiter(username, self.room, self.db)
|
|
if waiter is None:
|
|
await self.direct_send(type="error", payload={"msg": "Nom d'utilisateur invalide ou indisponible"})
|
|
return
|
|
self.member = waiter
|
|
self.add_to_group()
|
|
await self.direct_send(type="waiting", payload={"waiter": serialize_member(self.member).dict()})
|
|
await self.send_to_admin(type="waiter", payload={"waiter": serialize_member(self.member).dict()})
|
|
else:
|
|
if token is not None:
|
|
user = get_user_from_token(token, self.db)
|
|
if user is None:
|
|
await self.direct_send(type="error", payload={"msg": "Utilisateur introuvable"})
|
|
return
|
|
if user is False:
|
|
await self.direct_send(type="error", payload={"msg": "Token expired"})
|
|
return
|
|
|
|
member = create_user_member(user, self.room, self.db)
|
|
if member is None:
|
|
return
|
|
self.member = member
|
|
self.add_to_group()
|
|
await self.broadcast(type="joined", payload={"member": serialize_member(self.member).dict()}, exclude=True)
|
|
await self.direct_send(type="accepted", payload={"member": serialize_member(self.member).dict()})
|
|
elif username is not None:
|
|
member = create_anonymous_member(username, self.room, self.db)
|
|
if member is None:
|
|
await self.direct_send(type="error", data={"msg": "Nom d'utilisateur indisponible"})
|
|
return
|
|
self.member = member
|
|
self.add_to_group()
|
|
|
|
await self.broadcast(type="joined", payload={"member": serialize_member(self.member).dict()}, exclude=True)
|
|
await self.direct_send(type="accepted", payload={"member": serialize_member(self.member).dict()})
|
|
|
|
async def isAdminReceive(self):
|
|
is_admin = self.member is not None and self.member.is_admin == True
|
|
if not is_admin:
|
|
await self.direct_send(type="error", payload={"msg": "Vous n'avez pas la permission de faire ca"})
|
|
return False
|
|
return True
|
|
|
|
def isAdmin(self):
|
|
return self.member is not None and self.member.is_admin == True
|
|
|
|
@Consumer.event('accept', conditions=[isAdminReceive])
|
|
async def accept(self, waiter_id: str):
|
|
waiter = get_waiter(waiter_id, self.db)
|
|
if waiter is None:
|
|
await self.direct_send(type="error", payload={'msg': "Utilisateur introuvable"})
|
|
return
|
|
member = accept_waiter(waiter, self.db)
|
|
await self.send_to(type="accepted", payload={"member": serialize_member(member).dict()}, member_id=waiter_id)
|
|
await self.broadcast(type="joined", payload={"member": serialize_member(member).dict()})
|
|
|
|
@Consumer.event('refuse', conditions=[isAdminReceive])
|
|
async def accept(self, waiter_id: str):
|
|
waiter = get_waiter(waiter_id, self.db)
|
|
member = refuse_waiter(waiter, self.db)
|
|
await self.send_to(type="refused", payload={'waiter_id': waiter_id}, member_id=waiter_id)
|
|
# await self.broadcast(type="joined", payload={"member": serialize_member(member).dict()})
|
|
|
|
@Consumer.event('ping_room')
|
|
async def proom(self):
|
|
await self.broadcast(type='ping', payload={}, exclude=True)
|
|
async def hasRoom(self):
|
|
if self.member is None:
|
|
await self.direct_send(type="error", payload={"msg": "Vous n'êtes connecté à aucune salle"})
|
|
return self.member
|
|
@Consumer.event('leave', conditions=[hasRoom])
|
|
async def leave(self):
|
|
if self.member.is_admin is True:
|
|
await self.direct_send(type="error", payload={"msg": "Vous ne pouvez pas quitter une salle dont vous êtes l'administrateur"})
|
|
return
|
|
member_obj = serialize_member(self.member).dict()
|
|
leave_room(self.member, self.db)
|
|
await self.direct_send(type="successfully_leaved", payload = {})
|
|
await self.broadcast(type='leaved', payload ={"member": member_obj})
|
|
self.member = None
|
|
|
|
@Consumer.event('ban', conditions=[isAdminReceive])
|
|
async def ban(self, member_id: str):
|
|
member = get_member(member_id, self.room.id, self.db)
|
|
if member == None:
|
|
await self.direct_send(type="error", payload={"msg": "Utilisateur introuvable"})
|
|
return
|
|
if member.is_admin is True:
|
|
await self.direct_send(type="error", payload={"msg": "Vous ne pouvez pas bannir un administrateur"})
|
|
return
|
|
member_serialized = serialize_member(member)
|
|
|
|
leave_room(member, self.db)
|
|
await self.send_to(type="banned", payload={}, member_id=member.id_code)
|
|
await self.broadcast(type="leaved", payload = {"member": member_serialized.dict()})
|
|
|
|
def isMember(self):
|
|
return self.member is not None and self.member.waiting == False
|
|
|
|
# Sending Events
|
|
@Consumer.sending("joined", conditions=[isMember])
|
|
def joined(self, member: MemberRead):
|
|
if self.member.id_code == member.id_code:
|
|
raise ValueError("")
|
|
if self.member.is_admin == False:
|
|
member.reconnect_code = ""
|
|
return {"member": member}
|
|
|
|
@Consumer.sending(['waiter', "disconnect_waiter"], conditions=[isAdmin])
|
|
def waiter(self, waiter: Waiter):
|
|
return {"waiter": waiter}
|
|
|
|
@Consumer.sending(['connect', "disconnect", ""])
|
|
def connect_event(self, member: MemberRead):
|
|
if not self.member.is_admin:
|
|
member.reconnect_code = ""
|
|
return {"member": member}
|
|
|
|
@Consumer.sending('disconnect')
|
|
def disconnect_event(self, member: MemberRead):
|
|
if not self.member.is_admin:
|
|
member.reconnect_code = ""
|
|
return {"member": member}
|
|
|
|
@Consumer.sending('disconnect_waiter', conditions=[isAdmin])
|
|
def disconnect_event(self, waiter: Waiter):
|
|
return {"waiter": waiter}
|
|
def isWaiter(self):
|
|
return self.member is not None and self.member.waiting == True
|
|
|
|
@Consumer.sending("refused", conditions=[isWaiter])
|
|
def refused(self, waiter_id: str):
|
|
self.member = None
|
|
self.manager.remove(self.room.id, self)
|
|
return {"waiter_id": waiter_id}
|
|
|
|
@Consumer.sending("banned", conditions=[isMember])
|
|
def banned(self):
|
|
self.member = None
|
|
self.manager.remove(self.room.id, self)
|
|
return {}
|
|
|
|
@Consumer.sending('ping', conditions=[isMember])
|
|
def ping(self):
|
|
return {}
|
|
async def disconnect(self):
|
|
print("DISCONNECT", self.member)
|
|
self.manager.remove(self.room.id, self)
|
|
await self.disconnect_self()
|
|
|
|
|
|
|
|
class RoomManager:
|
|
def __init__(self):
|
|
self.active_connections: Dict[str, List[RoomConsumer]] = {}
|
|
|
|
def add(self, group: str, ws: RoomConsumer):
|
|
|
|
if group not in self.active_connections:
|
|
self.active_connections[group] = []
|
|
#print("adding", ws, self.active_connections[group])
|
|
if ws not in self.active_connections[group]:
|
|
#print('ACTUALLY ADDING')
|
|
self.active_connections[group].append(ws)
|
|
|
|
def remove(self, group: str, ws: RoomConsumer):
|
|
if group in self.active_connections:
|
|
if ws in self.active_connections[group]:
|
|
self.active_connections[group].remove(ws)
|
|
|
|
async def broadcast(self, message, group: str, exclude: list[RoomConsumer] = []):
|
|
if group in self.active_connections:
|
|
#print(self.active_connections[group], exclude)
|
|
for connection in list(set(self.active_connections[group])):
|
|
if connection not in exclude:
|
|
#print('SEND TO', connection, message)
|
|
await connection.send(message)
|
|
|
|
async def send_to(self, group, id_code, msg):
|
|
if group in self.active_connections:
|
|
members = [c for c in self.active_connections[group] if c.member.id_code == id_code]
|
|
for m in members:
|
|
await m.send(msg)
|
|
|
|
async def send_to_admin(self, group, msg):
|
|
if group in self.active_connections:
|
|
members = [c for c in self.active_connections[group] if c.member.is_admin == True]
|
|
for m in members:
|
|
await m.send(msg)
|
|
|
|
manager = RoomManager()
|
|
|
|
@router.websocket('/ws/room/{room_id}')
|
|
async def room_ws(ws: WebSocket, room: Room | None = Depends(check_room), db: Session = Depends(get_session)):
|
|
if room is None:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND, detail='Room not found')
|
|
consumer = RoomConsumer(ws=ws, room=room, manager=manager, db=db)
|
|
await consumer.run()
|