Update main.py

main.py

@@ -1,61 +1,51 @@
-from flask import Flask, request, jsonify, Response
+from flask import Flask, request, jsonify
 import random
 import string
-import threading
-from queue import Queue
+from threading import Lock
 
 app = Flask(__name__)
-transfer_queues = {}
-lock = threading.Lock()
+lock = Lock()
+transfers = {}
 
 def generate_id():
     return ''.join(random.choices(string.ascii_letters + string.digits, k=8))
 
 @app.route('/create_transfer', methods=['POST'])
 def create_transfer():
-    transfer_id = generate_id()
     with lock:
-        transfer_queues[transfer_id] = {
-            'queue': Queue(maxsize=10),
+        transfer_id = generate_id()
+        transfers[transfer_id] = {
+            'sender': None,
+            'receiver': None,
             'filename': request.json['filename'],
-            'active': False
+            'filesize': request.json['filesize']
         }
-    return jsonify({'transfer_id': transfer_id})
+        return jsonify({'transfer_id': transfer_id})
 
-@app.route('/upload/<transfer_id>', methods=['POST'])
-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('/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'})
 
-@app.route('/download/<transfer_id>')
-def download(transfer_id):
-    def generate():
-        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
-            except:
-                transfer_queues.pop(transfer_id, None)
+@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'})
 
-    headers = {
-        'Content-Disposition': f'attachment; filename="{transfer_queues[transfer_id]["filename"]}"'
-    }
-    return Response(generate(), mimetype='application/octet-stream', headers=headers)
+@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')
+    })
 
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=5000, threaded=True)
+    app.run(host='0.0.0.0', port=5000)