|
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 |
|
CHUNK_SIZE = 100 * 1024 * 1024 |
|
MAX_FILE_SIZE = 20 * 1024 * 1024 * 1024 |
|
|
|
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 |
|
|
|
|
|
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) |