Update main.py

main.py

@@ -1,67 +1,61 @@
 from flask import Flask, request, jsonify, Response
 import random
 import string
-from collections import deque
 import threading
+from queue import Queue
 
 app = Flask(__name__)
+transfer_queues = {}
+lock = threading.Lock()
 
-transfers = {}
-transfer_data = {}
-transfer_lock = threading.Lock()
-
-def generate_short_id(length=4):
-    while True:
-        token = ''.join(random.choices(string.ascii_letters + string.digits, k=length))
-        if token not in transfers:
-            return token
+def generate_id():
+    return ''.join(random.choices(string.ascii_letters + string.digits, k=8))
 
 @app.route('/create_transfer', methods=['POST'])
 def create_transfer():
-    filename = request.json.get('filename')
-    filesize = request.json.get('filesize', 0)
-    
-    transfer_id = generate_short_id(4)
-    with transfer_lock:
-        transfers[transfer_id] = {
-            'filename': filename,
-            'filesize': filesize,
-            'completed': False
+    transfer_id = generate_id()
+    with lock:
+        transfer_queues[transfer_id] = {
+            'queue': Queue(maxsize=10),
+            'filename': request.json['filename'],
+            'active': False
         }
-        transfer_data[transfer_id] = deque()
     return jsonify({'transfer_id': transfer_id})
 
 @app.route('/upload/<transfer_id>', methods=['POST'])
-def upload_file(transfer_id):
-    if transfer_id not in transfers:
-        return jsonify({'error': 'Invalid transfer ID'}), 404
-    
-    data = request.data
-    with transfer_lock:
-        transfer_data[transfer_id].append(data)
-        transfers[transfer_id]['completed'] = request.headers.get('X-Transfer-Complete') == 'true'
-    
-    return jsonify({'status': 'chunk uploaded'})
-
-@app.route('/download/<transfer_id>', methods=['GET'])
-def download_file(transfer_id):
-    if transfer_id not in transfers:
-        return jsonify({'error': 'Invalid transfer ID'}), 404
-
+def upload(transfer_id):
+    if transfer_id not in transfer_queues:
+        return jsonify({'error': 'Invalid ID'}), 404
+
+    chunk = request.data
+    try:
+        transfer_queues[transfer_id]['queue'].put(chunk, timeout=30)
+        return jsonify({'status': 'ok'})
+    except:
+        return jsonify({'error': 'Receiver not connected'}), 408
+
+@app.route('/download/<transfer_id>')
+def download(transfer_id):
     def generate():
-        while True:
-            with transfer_lock:
-                if transfer_data[transfer_id]:
-                    chunk = transfer_data[transfer_id].popleft()
+        with lock:
+            if transfer_id not in transfer_queues:
+                yield 'Transfer ID invalid'
+                return
+            
+            transfer = transfer_queues[transfer_id]
+            transfer['active'] = True
+            
+            try:
+                while True:
+                    chunk = transfer['queue'].get(timeout=30)
                     yield chunk
-                elif transfers[transfer_id]['completed']:
-                    break
+            except:
+                transfer_queues.pop(transfer_id, None)
 
-    return Response(
-        generate(),
-        mimetype='application/octet-stream',
-        headers={'Content-Disposition': f'attachment; filename="{transfers[transfer_id]["filename"]}"'}
-    )
+    headers = {
+        'Content-Disposition': f'attachment; filename="{transfer_queues[transfer_id]["filename"]}"'
+    }
+    return Response(generate(), mimetype='application/octet-stream', headers=headers)
 
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=5000)
+    app.run(host='0.0.0.0', port=5000, threaded=True)