Update main.py

main.py

@@ -1,87 +1,109 @@
-from flask import Flask, request, Response
+from flask import Flask, request, jsonify, Response
 import random
 import string
-from collections import deque
-import threading
 import time
+import threading
+import os
+from collections import defaultdict
 
 app = Flask(__name__)
 
-transfers = {}
-transfer_lock = threading.Lock()
-BUFFER_SIZE = 5  # Максимальное количество хранимых чанков
+# Конфигурация
+MAX_MEMORY = 12 * 1024 * 1024 * 1024  # 12 GB
+CHUNK_SIZE = 5 * 1024 * 1024  # 5 MB
+FILE_STORAGE = "./temp_uploads"
+os.makedirs(FILE_STORAGE, exist_ok=True)
 
-class TransferSession:
-    def __init__(self, filename):
-        self.filename = filename
-        self.buffer = deque(maxlen=BUFFER_SIZE)
-        self.completed = False
-        self.event = threading.Event()
-        self.last_activity = time.time()
+active_transfers = defaultdict(dict)
+transfer_lock = threading.Lock()
 
-def generate_id():
-    return ''.join(random.choices(string.ascii_letters + string.digits, k=8))
+def generate_transfer_id():
+    return ''.join(random.choices(string.ascii_letters + string.digits, k=16))
 
-@app.route('/start_upload', methods=['POST'])
-def start_upload():
-    filename = request.json.get('filename', 'file.bin')
-    transfer_id = generate_id()
+@app.route('/init_upload', methods=['POST'])
+def init_upload():
+    filename = request.json.get('filename')
+    filesize = request.json.get('filesize')
+    
+    transfer_id = generate_transfer_id()
+    filepath = os.path.join(FILE_STORAGE, transfer_id)
     
     with transfer_lock:
-        transfers[transfer_id] = TransferSession(filename)
+        active_transfers[transfer_id] = {
+            'filename': filename,
+            'filesize': filesize,
+            'received': 0,
+            'filepath': filepath,
+            'created_at': time.time()
+        }
+    
+    # Создаем пустой файл
+    with open(filepath, 'wb') as f:
+        f.truncate(filesize)
     
-    return {'transfer_id': transfer_id}
+    return jsonify({'transfer_id': transfer_id})
 
 @app.route('/upload/<transfer_id>', methods=['POST'])
 def upload_chunk(transfer_id):
-    with transfer_lock:
-        session = transfers.get(transfer_id)
-        if not session:
-            return {'error': 'Invalid transfer ID'}, 404
-        
-        chunk = request.data
-        session.buffer.append(chunk)
-        session.last_activity = time.time()
-        session.event.set()
-        session.event.clear()
-        
-        if request.headers.get('X-Transfer-Complete') == 'true':
-            session.completed = True
+    if transfer_id not in active_transfers:
+        return jsonify({'error': 'Invalid transfer ID'}), 404
+    
+    chunk = request.data
+    start = int(request.headers.get('Content-Range').split('-')[0])
+    
+    filepath = active_transfers[transfer_id]['filepath']
+    
+    try:
+        with open(filepath, 'r+b') as f:
+            f.seek(start)
+            f.write(chunk)
         
-        return {'status': 'chunk accepted'}
-
-@app.route('/download/<transfer_id>')
-def download_chunks(transfer_id):
-    def generate():
         with transfer_lock:
-            session = transfers.get(transfer_id)
-            if not session:
-                yield b'Transfer not found'
-                return
+            active_transfers[transfer_id]['received'] += len(chunk)
             
+        return jsonify({'status': 'success', 'received': active_transfers[transfer_id]['received']})
+    
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+@app.route('/download/<transfer_id>', methods=['GET'])
+def download_file(transfer_id):
+    def generate():
+        filepath = active_transfers[transfer_id]['filepath']
+        with open(filepath, 'rb') as f:
             while True:
-                while session.buffer:
-                    chunk = session.buffer.popleft()
-                    yield chunk
-                
-                if session.completed:
-                    break
-                
-                session.event.wait(timeout=30)
-                if time.time() - session.last_activity > 60:
+                chunk = f.read(CHUNK_SIZE)
+                if not chunk:
                     break
-
-        # Очистка через минуту после завершения
-        time.sleep(60)
-        with transfer_lock:
-            if transfer_id in transfers:
-                del transfers[transfer_id]
-
+                yield chunk
+        
+        # Удаляем файл после отправки
+        os.remove(filepath)
+        del active_transfers[transfer_id]
+    
     return Response(
         generate(),
         mimetype='application/octet-stream',
-        headers={'Content-Disposition': f'attachment; filename="{session.filename}"'}
+        headers={
+            'Content-Disposition': f'attachment; filename="{active_transfers[transfer_id]["filename"]}"',
+            'Content-Length': str(active_transfers[transfer_id]['filesize'])
+        }
     )
 
+# Очистка зависших передач
+def cleanup():
+    while True:
+        time.sleep(3600)
+        now = time.time()
+        for tid in list(active_transfers.keys()):
+            if now - active_transfers[tid]['created_at'] > 86400:  # 24 часа
+                try:
+                    os.remove(active_transfers[tid]['filepath'])
+                    del active_transfers[tid]
+                except:
+                    pass
+
+threading.Thread(target=cleanup, daemon=True).start()
+
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=5000, threaded=True)