|
from flask import Flask, request, jsonify, Response |
|
import random |
|
import string |
|
import time |
|
import threading |
|
import os |
|
from collections import defaultdict |
|
|
|
app = Flask(__name__) |
|
|
|
|
|
MAX_MEMORY = 12 * 1024 * 1024 * 1024 |
|
CHUNK_SIZE = 5 * 1024 * 1024 |
|
FILE_STORAGE = "./temp_uploads" |
|
os.makedirs(FILE_STORAGE, exist_ok=True) |
|
|
|
active_transfers = defaultdict(dict) |
|
transfer_lock = threading.Lock() |
|
|
|
def generate_transfer_id(): |
|
return ''.join(random.choices(string.ascii_letters + string.digits, k=16)) |
|
|
|
@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: |
|
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 jsonify({'transfer_id': transfer_id}) |
|
|
|
@app.route('/upload/<transfer_id>', methods=['POST']) |
|
def upload_chunk(transfer_id): |
|
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) |
|
|
|
with transfer_lock: |
|
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: |
|
chunk = f.read(CHUNK_SIZE) |
|
if not chunk: |
|
break |
|
yield chunk |
|
|
|
|
|
os.remove(filepath) |
|
del active_transfers[transfer_id] |
|
|
|
return Response( |
|
generate(), |
|
mimetype='application/octet-stream', |
|
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: |
|
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) |