generateur_v3/backend/api_old/apis/room/websocket.py

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()