share / main.py
Starchik1's picture
Update main.py
146248c verified
raw
history blame
3.85 kB
from flask import Flask, request, jsonify, Response
import random
import string
import os
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=16))
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
manager = TransferManager()
@app.route('/init_upload', methods=['POST'])
def init_upload():
filename = request.json.get('filename')
filesize = request.json.get('filesize')
try:
transfer_id = manager.create_transfer(filename, filesize)
return jsonify({'transfer_id': transfer_id})
except Exception as e:
return jsonify({'error': str(e)}), 400
@app.route('/upload/<transfer_id>', methods=['POST'])
def upload_chunk(transfer_id):
if transfer_id not in manager.transfers:
return jsonify({'error': 'Invalid transfer ID'}), 404
chunk = request.data
chunk_size = len(chunk)
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
@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
with open(file_path, 'rb') as f:
while True:
chunk = f.read(CHUNK_SIZE)
if not chunk:
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"]}"'}
)
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)