Update main.py

main.py

@@ -1,122 +1,67 @@
 from flask import Flask, request, jsonify, Response
 import random
 import string
-import os
+from collections import deque
 import threading
-import time
-from pathlib import Path
 
 app = Flask(__name__)
-app.config['UPLOAD_FOLDER'] = 'temp_uploads'
-Path(app.config['UPLOAD_FOLDER']).mkdir(exist_ok=True)
-
-# Конфигурация памяти
-MAX_MEMORY = 12 * 1024 * 1024 * 1024  # 12 GB
-CHUNK_SIZE = 100 * 1024 * 1024  # 4 MB
-MAX_FILE_SIZE = 20 * 1024 * 1024 * 1024  # 20 GB
 
 transfers = {}
-active_memory = 0
-lock = threading.Lock()
-
-class TransferManager:
-    def __init__(self):
-        self.transfers = {}
-        self.file_locks = {}
-    
-    def create_transfer(self, filename, filesize):
-        if filesize > MAX_FILE_SIZE:
-            raise ValueError("File size exceeds maximum limit")
-            
-        transfer_id = ''.join(random.choices(string.ascii_letters + string.digits, k=4))
-        file_path = os.path.join(app.config['UPLOAD_FOLDER'], transfer_id)
-        
-        with lock:
-            self.transfers[transfer_id] = {
-                'filename': filename,
-                'filesize': filesize,
-                'received': 0,
-                'file_path': file_path,
-                'created_at': time.time(),
-                'status': 'uploading'
-            }
-            self.file_locks[transfer_id] = threading.Lock()
-        
-        return transfer_id
+transfer_data = {}
+transfer_lock = threading.Lock()
 
-manager = TransferManager()
+def generate_short_id(length=8):
+    while True:
+        token = ''.join(random.choices(string.ascii_letters + string.digits, k=length))
+        if token not in transfers:
+            return token
 
-@app.route('/init_upload', methods=['POST'])
-def init_upload():
+@app.route('/create_transfer', methods=['POST'])
+def create_transfer():
     filename = request.json.get('filename')
-    filesize = request.json.get('filesize')
+    filesize = request.json.get('filesize', 0)
     
-    try:
-        transfer_id = manager.create_transfer(filename, filesize)
-        return jsonify({'transfer_id': transfer_id})
-    except Exception as e:
-        return jsonify({'error': str(e)}), 400
+    transfer_id = generate_short_id(8)
+    with transfer_lock:
+        transfers[transfer_id] = {
+            'filename': filename,
+            'filesize': filesize,
+            'completed': False
+        }
+        transfer_data[transfer_id] = deque()
+    return jsonify({'transfer_id': transfer_id})
 
 @app.route('/upload/<transfer_id>', methods=['POST'])
-def upload_chunk(transfer_id):
-    if transfer_id not in manager.transfers:
+def upload_file(transfer_id):
+    if transfer_id not in transfers:
         return jsonify({'error': 'Invalid transfer ID'}), 404
     
-    chunk = request.data
-    chunk_size = len(chunk)
+    data = request.data
+    with transfer_lock:
+        transfer_data[transfer_id].append(data)
+        transfers[transfer_id]['completed'] = request.headers.get('X-Transfer-Complete') == 'true'
     
-    with lock:
-        transfer = manager.transfers[transfer_id]
-        if transfer['status'] != 'uploading':
-            return jsonify({'error': 'Transfer completed'}), 400
-    
-    try:
-        with manager.file_locks[transfer_id], open(transfer['file_path'], 'ab') as f:
-            f.write(chunk)
-            transfer['received'] += chunk_size
-        
-        if transfer['received'] >= transfer['filesize']:
-            transfer['status'] = 'ready'
-        
-        return jsonify({
-            'received': transfer['received'],
-            'status': transfer['status']
-        })
-    except Exception as e:
-        return jsonify({'error': str(e)}), 500
+    return jsonify({'status': 'chunk uploaded'})
 
 @app.route('/download/<transfer_id>', methods=['GET'])
 def download_file(transfer_id):
-    def generate():
-        file_path = manager.transfers.get(transfer_id, {}).get('file_path')
-        if not file_path or not os.path.exists(file_path):
-            yield b'File not found'
-            return
+    if transfer_id not in transfers:
+        return jsonify({'error': 'Invalid transfer ID'}), 404
 
-        with open(file_path, 'rb') as f:
-            while True:
-                chunk = f.read(CHUNK_SIZE)
-                if not chunk:
+    def generate():
+        while True:
+            with transfer_lock:
+                if transfer_data[transfer_id]:
+                    chunk = transfer_data[transfer_id].popleft()
+                    yield chunk
+                elif transfers[transfer_id]['completed']:
                     break
-                yield chunk
-        
-        # Автоочистка через 1 час
-        threading.Timer(3600, cleanup_transfer, args=[transfer_id]).start()
-    
+
     return Response(
         generate(),
         mimetype='application/octet-stream',
-        headers={'Content-Disposition': f'attachment; filename="{manager.transfers[transfer_id]["filename"]}"'}
+        headers={'Content-Disposition': f'attachment; filename="{transfers[transfer_id]["filename"]}"'}
     )
 
-def cleanup_transfer(transfer_id):
-    with lock:
-        if transfer_id in manager.transfers:
-            try:
-                os.remove(manager.transfers[transfer_id]['file_path'])
-            except:
-                pass
-            del manager.transfers[transfer_id]
-
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=5000, threaded=True)
+    app.run(host='0.0.0.0', port=5000)