share / main.py
Starchik1's picture
Update main.py
5c18c00 verified
raw
history blame
3.79 kB
from flask import Flask, request, jsonify, Response
import random
import string
import threading
import time
import os
from pathlib import Path
app = Flask(__name__)
app.config['MAX_CONTENT_LENGTH'] = 12 * 1024 * 1024 * 1024 # 12GB
UPLOAD_FOLDER = Path('uploads')
UPLOAD_FOLDER.mkdir(exist_ok=True)
transfers = {}
transfer_lock = threading.Lock()
TIMEOUT = 3600 # 1 час
CHUNK_SIZE = 1024 * 1024 # 1MB
def generate_short_id(length=10):
while True:
token = ''.join(random.choices(string.ascii_letters + string.digits, k=length))
if token not in transfers:
return token
@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()
file_path = UPLOAD_FOLDER / f"{transfer_id}.dat"
with transfer_lock:
transfers[transfer_id] = {
'filename': filename,
'filesize': filesize,
'completed': False,
'last_activity': time.time(),
'file_path': file_path,
'file_handle': open(file_path, 'wb')
}
return jsonify({'transfer_id': transfer_id})
@app.route('/upload/<transfer_id>', methods=['POST'])
def upload_file(transfer_id):
with transfer_lock:
if transfer_id not in transfers:
return jsonify({'error': 'Invalid transfer ID'}), 404
transfer = transfers[transfer_id]
transfer['last_activity'] = time.time()
try:
# Записываем полученные данные в файл
data = request.data
transfer['file_handle'].write(data)
transfer['file_handle'].flush()
if request.headers.get('X-Transfer-Complete') == 'true':
transfer['completed'] = True
transfer['file_handle'].close()
return jsonify({'status': 'chunk uploaded', 'received': len(data)})
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 = None
with transfer_lock:
if transfer_id not in transfers:
yield b'Transfer expired'
return
transfer = transfers[transfer_id]
file_path = transfer['file_path']
transfer['last_activity'] = time.time()
if not file_path.exists():
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
return Response(
generate(),
mimetype='application/octet-stream',
headers={'Content-Disposition': f'attachment; filename="{transfers[transfer_id]["filename"]}"'}
)
def cleanup_thread():
while True:
time.sleep(60 * 5) # Проверка каждые 5 минут
with transfer_lock:
now = time.time()
expired = []
for tid, t in transfers.items():
if now - t['last_activity'] > TIMEOUT:
expired.append(tid)
if 'file_handle' in t and not t['file_handle'].closed:
t['file_handle'].close()
if t['file_path'].exists():
t['file_path'].unlink()
for tid in expired:
del transfers[tid]
threading.Thread(target=cleanup_thread, daemon=True).start()
if __name__ == '__main__':