from typing import List from fastapi import WebSocket, WebSocketException class ConnectionManager: def __init__(self, password: str) -> None: self.__password: str = password self.__active_connections: List[WebSocket] = list[WebSocket]() async def connect(self, websocket: WebSocket) -> None: authorization: str = websocket.headers.get("Authorization") if not authorization: raise WebSocketException(code=1008, reason="Authorization header missing") token_type, _, token = authorization.partition(' ') if token_type != "Bearer" or not token: raise WebSocketException(code=1008, reason="Invalid authorization header format") if token != self.__password: raise WebSocketException(code=1008, reason="Invalid token") await websocket.accept() self.__active_connections.append(websocket) def disconnect(self, websocket: WebSocket) -> None: self.__active_connections.remove(websocket)