619 lines
23 KiB
Python
619 lines
23 KiB
Python
from pydantic import Field
|
|
from pydantic import create_model_from_typeddict
|
|
from pydantic.error_wrappers import ValidationError
|
|
from pydantic import validate_arguments
|
|
from services.database import generate_unique_code
|
|
from sqlmodel import col
|
|
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 database.room.crud import check_room, create_room_db, userInRoom
|
|
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)
|
|
if member.waiting == True:
|
|
return Waiter(username=member_obj.username, waiter_id=member.waiter_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, conditions: List[Callable | bool] = []):
|
|
def add_event(func):
|
|
model = validate_arguments(func).model
|
|
eventsDict[name] = {"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.id)
|
|
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:
|
|
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:
|
|
|
|
print('NOPE', self.member.id, 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)
|
|
|
|
|
|
manager = ConsumerManager()
|
|
|
|
|
|
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 get_member_from_user(user_id: int, room_id: int, db: Session):
|
|
member = db.exec(select(Member).where(Member.room_id ==
|
|
room_id, Member.user_id == user_id)).first()
|
|
return member
|
|
|
|
|
|
def get_member_from_token(token: str, room_id: int, db: Session):
|
|
user = get_user_from_token(token, db)
|
|
if user is False:
|
|
return False
|
|
if user is None:
|
|
return None
|
|
member = get_member_from_user(user.id, room_id, db)
|
|
return member
|
|
|
|
|
|
def get_member_from_anonymous(anonymous_id: int, room_id: int, db: Session):
|
|
member = db.exec(select(Member).where(Member.room_id ==
|
|
room_id, Member.anonymous_id == anonymous_id)).first()
|
|
return member
|
|
|
|
|
|
def get_member_reconnect_code(reconnect_code: str, room_id: int, db: Session):
|
|
anonymous = get_anonymous_from_code(reconnect_code, db)
|
|
if anonymous is None:
|
|
return None
|
|
member = get_member_from_anonymous(anonymous.id, room_id, db)
|
|
return member
|
|
|
|
|
|
def get_anonymous_from_code(reconnect_code: str, db: Session):
|
|
anonymous = db.exec(select(Anonymous).where(
|
|
Anonymous.reconnect_code == reconnect_code)).first()
|
|
return anonymous
|
|
|
|
|
|
def connect_member(member: Member, db: Session):
|
|
member.online = True
|
|
db.add(member)
|
|
db.commit()
|
|
db.refresh(member)
|
|
return member
|
|
|
|
|
|
def disconnect_member(member: Member, db: Session):
|
|
if member.waiting == False:
|
|
member.online = False
|
|
db.add(member)
|
|
db.commit()
|
|
db.refresh(member)
|
|
return member
|
|
else:
|
|
db.delete(member)
|
|
db.commit()
|
|
return member
|
|
|
|
|
|
def validate_username(username: str, room: Room, db: Session = Depends(get_session)):
|
|
print('VALIDATE', username)
|
|
if len(username) > 20:
|
|
return None
|
|
members = select(Member.anonymous_id).where(
|
|
Member.room_id == room.id, Member.anonymous_id != None)
|
|
anonymous = select(Anonymous).where(
|
|
col(Anonymous.id).in_(members), Anonymous.username == username)
|
|
username_anonymous = db.exec(anonymous).first()
|
|
return None if username_anonymous is not None else username
|
|
|
|
|
|
def create_anonymous_member(username: str, room: Room, db: Session):
|
|
username = validate_username(username)
|
|
if username is None:
|
|
return None
|
|
reconnect_code = generate_unique_code(
|
|
Anonymous, s=db, field_name="reconnect_code")
|
|
anonymous = Anonymous(username=username, reconnect_code=reconnect_code)
|
|
member = Member(room=room, anonymous=anonymous)
|
|
db.add(member)
|
|
db.commit()
|
|
db.refresh(member)
|
|
return member
|
|
|
|
|
|
def create_user_member(user: User, room: Room, db: Session):
|
|
member = get_member_from_user(user.id, room.id, db)
|
|
if member is not None:
|
|
return None
|
|
member = Member(room=room, user=user)
|
|
db.add(member)
|
|
db.commit()
|
|
db.refresh(member)
|
|
return member
|
|
|
|
|
|
def create_anonymous_waiter(username: str, room: Room, db: Session):
|
|
username = validate_username(username, room, db)
|
|
if username is None:
|
|
return None
|
|
reconnect_code = generate_unique_code(
|
|
Anonymous, s=db, field_name="reconnect_code")
|
|
anonymous = Anonymous(username=username, reconnect_code=reconnect_code)
|
|
|
|
waiter_code = generate_unique_code(Member, s=db, field_name="waiter_code")
|
|
member = Member(room=room, anonymous=anonymous,
|
|
waiting=True, waiter_code=waiter_code)
|
|
db.add(member)
|
|
db.commit()
|
|
db.refresh(member)
|
|
return member
|
|
|
|
|
|
def create_user_waiter(user: User, room: Room, db: Session):
|
|
member = get_member_from_user(user.id, room.id, db)
|
|
if member is not None:
|
|
return None
|
|
waiter_code = generate_unique_code(Member, s=db, field_name="waiter_code")
|
|
member = Member(room=room, user=user, waiting=True,
|
|
waiter_code=waiter_code)
|
|
db.add(member)
|
|
db.commit()
|
|
db.refresh(member)
|
|
return member
|
|
|
|
|
|
def get_waiter(waiter_code: str, db: Session):
|
|
return db.exec(select(Member).where(Member.waiter_code == waiter_code)).first()
|
|
|
|
|
|
def accept_waiter(member: Member, db: Session):
|
|
member.waiting = False
|
|
member.waiter_code = None
|
|
db.add(member)
|
|
db.commit()
|
|
db.refresh(member)
|
|
return member
|
|
|
|
def refuse_waiter(member: Member, db: Session):
|
|
db.delete(member)
|
|
db.commit()
|
|
return None
|
|
|
|
|
|
class RoomConsumer(Consumer):
|
|
|
|
def __init__(self, ws: WebSocket, room: Room, manager: ConsumerManager, 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_all_room(self, type: str, payload: Any, exclude: bool = False):
|
|
await self.manager.broadcast({'type': type, "data": payload}, f'{self.room.id}__member', [exclude == True and self.ws])
|
|
await self.manager.broadcast({'type': type, "data": payload}, f'{self.room.id}__admin', [exclude == True and self.ws])
|
|
|
|
async def send_to_admin(self, type: str, payload: Any, exclude: bool = False):
|
|
await self.manager.broadcast({'type': type, "data": payload}, f'{self.room.id}__admin', [exclude == True and self.ws])
|
|
|
|
async def send_to_members(self, type: str, payload: Any, exclude: bool = False):
|
|
await self.manager.broadcast({'type': type, "data": payload}, f'{self.room.id}__member', [exclude == True and self.ws])
|
|
|
|
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_admin(self):
|
|
self.manager.add(f'{self.room.id}__admin', self.ws)
|
|
|
|
def add_to_members(self):
|
|
self.manager.add(f'{self.room.id}__members', self.ws)
|
|
|
|
def add_to_group(self):
|
|
if self.member.waiting == True:
|
|
self.manager.add(f'waiter__{self.member.waiter_code}', 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()})
|
|
|
|
async def disconnect_self(self):
|
|
if isinstance(self.member, Member):
|
|
disconnect_member(self.member, self.db)
|
|
await self.broadcast(type="disconnect", payload={"member": 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_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.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.broadcast(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.broadcast(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:
|
|
return
|
|
if user is False:
|
|
return
|
|
|
|
member = create_user_member(user, self.room, self.db)
|
|
if member is None:
|
|
return
|
|
self.member = member
|
|
self.add_to_group()
|
|
await self.direct_send()
|
|
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.direct_send()
|
|
|
|
def isAdmin(self):
|
|
return self.member is not None and self.member.is_admin == True
|
|
|
|
@Consumer.event('accept', conditions=[isAdmin])
|
|
async def accept(self, waiter_id: str):
|
|
waiter = get_waiter(waiter_id, self.db)
|
|
member = accept_waiter(waiter, self.db)
|
|
await self.manager.broadcast({"type": "accepted", "data": {'member': serialize_member(member).dict()}}, f"waiter__{waiter_id}")
|
|
await self.broadcast(type="joined", payload={"member": serialize_member(member).dict()})
|
|
|
|
@Consumer.event('refuse', conditions=[isAdmin])
|
|
async def accept(self, waiter_id: str):
|
|
waiter = get_waiter(waiter_id, self.db)
|
|
member = refuse_waiter(waiter, self.db)
|
|
await self.manager.broadcast({"type": "refused", "data": {'waiter_id': waiter_id}}, f"waiter__{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={})
|
|
|
|
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):
|
|
print('MEMBER', self.member, member)
|
|
if (self.member.user is not None and member.username == self.member.user.username) or (self.member.anonymous and member.reconnect_code == self.member.anonymous.reconnect_code):
|
|
raise ValueError("Nope")
|
|
if self.member.is_admin == False:
|
|
member.reconnect_code = ""
|
|
return {"member": member}
|
|
|
|
@Consumer.sending('waiter', conditions=[isAdmin])
|
|
def waiter(self, waiter: Waiter):
|
|
return {"waiter": waiter}
|
|
|
|
@Consumer.sending('ping', conditions=[isMember])
|
|
def ping(self):
|
|
return {}
|
|
|
|
async def disconnect(self):
|
|
self.manager.remove(self.room.id, self)
|
|
#await self.disconnect_self()
|
|
|
|
|
|
@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()
|
|
|
|
|
|
class TestConsumer(Consumer):
|
|
async def connect(self):
|
|
await self.ws.accept()
|
|
|
|
def test(self):
|
|
return True
|
|
|
|
@Consumer.event("test", conditions=[True, test])
|
|
async def testering(self):
|
|
await self.ws.send_json({"type": "success"})
|
|
await self.send({"type": "test", "data": {"i": {"username": "lilian", "reconnect_code": "something", "isAdmin": False, "isUser": False}, "test": 12}})
|
|
# await self.send({"type": "test", "data": {"i": {"username": "lilian", "reconnect_code": "something", "isAdmin": False, "isUser": False}}})
|
|
return
|
|
|
|
@Consumer.sending('test', conditions=[])
|
|
def sendtest(self, i: MemberRead, test: int):
|
|
print("i", i)
|
|
print(i.reconnect_code)
|
|
print(dict(i))
|
|
i.reconnect_code = "nope"
|
|
return {"i": i, "test": test}
|
|
|
|
|
|
@router.websocket('/ws/test')
|
|
async def test(ws: WebSocket):
|
|
consumer = TestConsumer(ws)
|
|
await consumer.run()
|