|
from flask import Flask, request, Response |
|
import random |
|
import string |
|
from collections import deque |
|
import threading |
|
import time |
|
|
|
app = Flask(__name__) |
|
|
|
transfers = {} |
|
transfer_lock = threading.Lock() |
|
BUFFER_SIZE = 5 |
|
|
|
class TransferSession: |
|
def __init__(self, filename): |
|
self.filename = filename |
|
self.buffer = deque(maxlen=BUFFER_SIZE) |
|
self.completed = False |
|
self.event = threading.Event() |
|
self.last_activity = time.time() |
|
|
|
def generate_id(): |
|
return ''.join(random.choices(string.ascii_letters + string.digits, k=8)) |
|
|
|
@app.route('/start_upload', methods=['POST']) |
|
def start_upload(): |
|
filename = request.json.get('filename', 'file.bin') |
|
transfer_id = generate_id() |
|
|
|
with transfer_lock: |
|
transfers[transfer_id] = TransferSession(filename) |
|
|
|
return {'transfer_id': transfer_id} |
|
|
|
@app.route('/upload/<transfer_id>', methods=['POST']) |
|
def upload_chunk(transfer_id): |
|
with transfer_lock: |
|
session = transfers.get(transfer_id) |
|
if not session: |
|
return {'error': 'Invalid transfer ID'}, 404 |
|
|
|
chunk = request.data |
|
session.buffer.append(chunk) |
|
session.last_activity = time.time() |
|
session.event.set() |
|
session.event.clear() |
|
|
|
if request.headers.get('X-Transfer-Complete') == 'true': |
|
session.completed = True |
|
|
|
return {'status': 'chunk accepted'} |
|
|
|
@app.route('/download/<transfer_id>') |
|
def download_chunks(transfer_id): |
|
def generate(): |
|
with transfer_lock: |
|
session = transfers.get(transfer_id) |
|
if not session: |
|
yield b'Transfer not found' |
|
return |
|
|
|
while True: |
|
while session.buffer: |
|
chunk = session.buffer.popleft() |
|
yield chunk |
|
|
|
if session.completed: |
|
break |
|
|
|
session.event.wait(timeout=30) |
|
if time.time() - session.last_activity > 60: |
|
break |
|
|
|
|
|
time.sleep(60) |
|
with transfer_lock: |
|
if transfer_id in transfers: |
|
del transfers[transfer_id] |
|
|
|
return Response( |
|
generate(), |
|
mimetype='application/octet-stream', |
|
headers={'Content-Disposition': f'attachment; filename="{session.filename}"'} |
|
) |
|
|
|
if __name__ == '__main__': |
|
app.run(host='0.0.0.0', port=5000, threaded=True) |