Update main.py

main.py

@@ -1,51 +1,40 @@
-from flask import Flask, request, jsonify
-import random
-import string
-from threading import Lock
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi.responses import HTMLResponse
+import uuid
 
-app = Flask(__name__)
-lock = Lock()
+app = FastAPI()
 transfers = {}
 
-def generate_id():
-    return ''.join(random.choices(string.ascii_letters + string.digits, k=8))
+class ConnectionManager:
+    def __init__(self):
+        self.active_connections = {}
 
-@app.route('/create_transfer', methods=['POST'])
-def create_transfer():
-    with lock:
-        transfer_id = generate_id()
-        transfers[transfer_id] = {
-            'sender': None,
-            'receiver': None,
-            'filename': request.json['filename'],
-            'filesize': request.json['filesize']
-        }
-        return jsonify({'transfer_id': transfer_id})
+    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)
 
-@app.route('/register_sender/<transfer_id>', methods=['POST'])
-def register_sender(transfer_id):
-    with lock:
-        transfers[transfer_id]['sender'] = {
-            'ip': request.remote_addr,
-            'port': request.json['port']
-        }
-        return jsonify({'status': 'sender registered'})
+    def disconnect(self, transfer_id: str, websocket: WebSocket):
+        self.active_connections[transfer_id].remove(websocket)
 
-@app.route('/register_receiver/<transfer_id>', methods=['POST'])
-def register_receiver(transfer_id):
-    with lock:
-        transfers[transfer_id]['receiver'] = {
-            'ip': request.remote_addr,
-            'port': request.json['port']
-        }
-        return jsonify({'status': 'receiver registered'})
+manager = ConnectionManager()
 
-@app.route('/get_peer_info/<transfer_id>')
-def get_peer_info(transfer_id):
-    return jsonify({
-        'sender': transfers.get(transfer_id, {}).get('sender'),
-        'receiver': transfers.get(transfer_id, {}).get('receiver')
-    })
+@app.post("/create_transfer")
+async def create_transfer():
+    transfer_id = str(uuid.uuid4())
+    transfers[transfer_id] = {"status": "waiting"}
+    return {"transfer_id": transfer_id}
 
-if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=5000)
+@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)