|
from fastapi import FastAPI, WebSocket, WebSocketDisconnect |
|
from fastapi.responses import HTMLResponse |
|
import uuid |
|
|
|
app = FastAPI() |
|
transfers = {} |
|
|
|
class ConnectionManager: |
|
def __init__(self): |
|
self.active_connections = {} |
|
|
|
async def connect(self, transfer_id: str, websocket: WebSocket): |
|
await websocket.accept() |
|
if transfer_id not in self.active_connections: |
|
self.active_connections[transfer_id] = [] |
|
self.active_connections[transfer_id].append(websocket) |
|
|
|
def disconnect(self, transfer_id: str, websocket: WebSocket): |
|
self.active_connections[transfer_id].remove(websocket) |
|
|
|
manager = ConnectionManager() |
|
|
|
@app.post("/create_transfer") |
|
async def create_transfer(): |
|
transfer_id = str(uuid.uuid4()) |
|
transfers[transfer_id] = {"status": "waiting"} |
|
return {"transfer_id": transfer_id} |
|
|
|
@app.websocket("/ws/{transfer_id}") |
|
async def websocket_endpoint(websocket: WebSocket, transfer_id: str): |
|
await manager.connect(transfer_id, websocket) |
|
try: |
|
while True: |
|
data = await websocket.receive_bytes() |
|
|
|
for connection in manager.active_connections.get(transfer_id, []): |
|
if connection != websocket: |
|
await connection.send_bytes(data) |
|
except WebSocketDisconnect: |
|
manager.disconnect(transfer_id, websocket) |