""" Blueprint: API Obfuscator. Seven endpoints: - POST /api/upload — upload one or more files -> {session_id} - GET /api/session_files/ — diagnostic list of session files - POST /api/cancel/ — request cooperative stop - GET /api/process_stream/ — SSE: process files, per-file progress - POST /api/process/ — process all session files -> {status:"done"} (legacy) - GET /api/download/ — download ZIP (with timestamp name) - GET /api/csv/ — download CSV separately """ import io import json import queue import threading import time import zipfile import logging from datetime import datetime, timedelta from flask import (Blueprint, current_app, request, send_file, jsonify, Response, stream_with_context) from drhider import obfuscate_files, LLMClient from drhider.obfuscator import LLMIncompleteError from upload.backend.session import (create_session, add_file, get_files, store_result, get_result, store_csv, get_csv, cleanup, file_count, MAX_FILE_BYTES, pause_ttl, resume_ttl, request_cancel, get_cancel_event) from upload.backend.upload_refs import safe_name api_bp = Blueprint("api", __name__, url_prefix="/api") log = logging.getLogger("routes.api_bp") def _disconnect_exceptions(): """Исключения, означающие отключение клиента SSE.""" return (GeneratorExit, BrokenPipeError, ConnectionResetError) @api_bp.route("/upload", methods=["POST"]) def upload(): """Upload files to session (один или несколько).""" sid = request.form.get("session", "") uploaded = request.files.getlist("files") if not uploaded: log.warning("upload: no files, sid=%s", sid) return jsonify({"ok": False, "error": "No file"}), 400 # CSV input is temporarily disabled so reliability testing can focus on the # other document formats. Keep the CSV processing code in the core intact. if any(safe_name(f.filename).lower().endswith(".csv") for f in uploaded): log.warning("upload: CSV input is disabled") return jsonify({"ok": False, "error": "CSV uploads are temporarily disabled"}), 400 if not sid: sid = create_session() added = 0 had_unnamed = False for f in uploaded: name = safe_name(f.filename) if not name: had_unnamed = True continue # Ранняя проверка размера по Content-Length позволяет не читать в память # тело файла, заведомо превышающего лимит. Проверка выполняется до f.read(). if f.content_length is not None and f.content_length > MAX_FILE_BYTES: log.warning("upload: file exceeds %dMB, skipped sid=%s file=%r size=%d", MAX_FILE_BYTES // (1024 * 1024), sid, name, f.content_length) continue data = f.read() log.info("upload: sid=%s file=%r size=%d", sid, name, len(data)) if len(data) > MAX_FILE_BYTES: log.warning("upload: file exceeds %dMB, skipped sid=%s file=%r size=%d", MAX_FILE_BYTES // (1024 * 1024), sid, name, len(data)) continue if not add_file(sid, name, data): log.warning("upload: session not found/limit, sid=%s file=%r", sid, f.filename) return jsonify({"ok": False, "error": "Session not found"}), 404 added += 1 if added == 0: err = "No filename" if had_unnamed else "No file" log.warning("upload: %s, sid=%s", err, sid) return jsonify({"ok": False, "error": err}), 400 log.info("upload: done sid=%s added=%d total=%d", sid, added, file_count(sid)) return jsonify({"ok": True, "session": sid, "count": file_count(sid)}) @api_bp.route("/session_files/", methods=["GET"]) def session_files(sid): """Отладка: список файлов сессии с размерами (для теста загрузки).""" if not current_app.config.get("ENABLE_SESSION_FILES"): return jsonify({"ok": False, "error": "Not found"}), 404 files = get_files(sid) if files is None: return jsonify({"ok": False, "error": "Session not found"}), 404 return jsonify({ "ok": True, "count": len(files), "files": [{"name": n, "size": len(b)} for n, b in files], }) @api_bp.route("/cancel/", methods=["POST"]) def cancel(sid): """Запросить мягкое прерывание обработки сессии. Воркер останавливается на ближайшей границе файла (текущий добирается), собирает частичный результат (готовые файлы + mapping) и шлёт SSE-событие `cancelled`. """ if request_cancel(sid): log.info("cancel: requested sid=%s", sid) return jsonify({"ok": True}), 200 log.warning("cancel: session not found sid=%s", sid) return jsonify({"ok": False, "error": "Session not found"}), 404 @api_bp.route("/process_stream/", methods=["GET"]) def process_stream(sid): """SSE: process all session files, streaming per-file progress. Все файлы обрабатываются ЕДИНЫМ вызовом obfuscate_files (общий mapping, согласованные токены). Обработка идёт в отдельном потоке; прогресс передаётся через очередь. Разрыв соединения клиента корректно перехватывается и останавливает генератор (и воркер — через cancel_event). """ files = get_files(sid) if files is None: return jsonify({"ok": False, "error": "Session not found"}), 404 if not files: return jsonify({"ok": False, "error": "No files"}), 400 all_files = [(fname, content, "") for fname, content in files] log.info("process_stream: start sid=%s files=%d", sid, len(all_files)) # A failed rerun must not leave an older ZIP or mapping downloadable. store_result(sid, None) store_csv(sid, None) # Сессия живёт, пока идёт обработка (TTL возобновляется в finally генератора) pause_ttl(sid) def generate(): llm = LLMClient() q = queue.Queue() cancel = threading.Event() # локальный: разрыв клиента (стоп heartbeat) cancel_event = get_cancel_event(sid) # из сессии: мягкая отмена (кнопка «Прервать») # Состояние для глобальной ETA (символы) eta = {"total_chars": 0, "done_chars": 0, "cur_chars": 0, "cur_total": 0, "cur_done": 0} _TOKENS_PER_CHAR = 8.0 # эмпирический коэффициент символов -> токенов LLM def progress(phase, idx, total_, name, elapsed, source_idxs): q.put(("progress", phase, idx, name, total_, elapsed, source_idxs)) # Состояние для per-file ETA (чанки) fstate = {"t0": 0.0, "total": 0, "done": 0} def file_progress(event, fname, **fields): if event == "file_start": fstate["t0"] = time.time() fstate["total"] = fields.get("chunks", 0) fstate["done"] = 0 elif event == "file_chunk": fstate["done"] = fields.get("chunks_done", 0) elapsed = time.time() - fstate["t0"] rate = fstate["done"] / elapsed if elapsed > 0 else 0 rem = fstate["total"] - fstate["done"] if rate > 0: fields["eta_sec"] = max(0, round(rem / rate)) q.put(("file", event, fname, fields)) def worker(): log.info("worker: start sid=%s files=%d", sid, len(all_files)) t0 = datetime.utcnow() try: zip_data, csv_str, meta = obfuscate_files( all_files, llm_client=llm, progress_cb=progress, cancel_event=cancel_event, file_progress=file_progress, ) stats = { "tokens": llm.tokens_total, "llm_sec": round(llm.llm_sec, 1), } dt = (datetime.utcnow() - t0).total_seconds() if meta and meta.get("cancelled"): log.info("worker: cancelled sid=%s in %.1fs processed=%d/%d", sid, dt, meta.get("processed", 0), meta.get("total", 0)) q.put(("cancelled", zip_data, csv_str, stats, meta)) else: log.info("worker: done sid=%s in %.1fs tokens=%d llm_sec=%.1f zip_len=%d", sid, dt, llm.tokens_total, llm.llm_sec, len(zip_data)) q.put(("result", zip_data, csv_str, stats)) except LLMIncompleteError as e: log.exception("worker: LLM incomplete sid=%s chunks=%d", sid, e.llm_errors) q.put(("error", { "error": "Processing stopped because the LLM could not process all fragments", "llm_incomplete": True, "llm_errors": e.llm_errors, })) except Exception: log.exception("worker: processing failed sid=%s", sid) q.put(("error", {"error": "Processing failed"})) threading.Thread(target=worker, daemon=True).start() log.debug("process_stream: worker thread started sid=%s", sid) try: while True: try: evt = q.get(timeout=1) except queue.Empty: if cancel.is_set(): log.info("process_stream: cancelled sid=%s (generator exit)", sid) return # Heartbeat: живая статистика LLM + глобальная ETA tokens = llm.tokens_total elapsed = llm.llm_elapsed_now() eta_sec = None if eta["total_chars"] > 0 and tokens > 0 and elapsed > 0: est_total_tokens = eta["total_chars"] * _TOKENS_PER_CHAR rate = tokens / elapsed if rate > 0: eta_sec = max(0, int((est_total_tokens - tokens) / rate)) done_chars = eta["done_chars"] if eta["cur_total"] > 0: done_chars += int(eta["cur_chars"] * (eta["cur_done"] / eta["cur_total"])) try: yield ( f"event: llm\n" f"data: {json.dumps({'active': llm.llm_active, 'elapsed': round(elapsed, 1), 'tokens': tokens, 'eta_sec': eta_sec, 'done_chars': done_chars, 'total_chars': eta['total_chars']})}\n\n" ) except _disconnect_exceptions() as e: log.warning("process_stream: disconnect during heartbeat sid=%s err=%r", sid, e) cancel.set() if cancel_event: cancel_event.set() return continue kind = evt[0] if kind == "progress": _, phase, idx, name, total_, elapsed, source_idxs = evt log.debug("process_stream: event=%s idx=%d name=%r elapsed=%s sid=%s", phase, idx, name, elapsed, sid) try: yield ( f"event: {phase}\n" f"data: {json.dumps({'idx': idx, 'source_idxs': source_idxs, 'name': name, 'total': total_, 'elapsed': elapsed})}\n\n" ) except _disconnect_exceptions() as e: log.warning("process_stream: disconnect on progress sid=%s phase=%s err=%r", sid, phase, e) cancel.set() if cancel_event: cancel_event.set() return elif kind == "file": _, event, fname, fields = evt # Обновляем состояние для глобальной ETA if event == "extract_complete": eta["total_chars"] = fields.get("total_chars", 0) elif event == "file_start": eta["cur_chars"] = fields.get("chars", 0) eta["cur_total"] = fields.get("chunks", 0) eta["cur_done"] = 0 elif event == "file_chunk": eta["cur_done"] = fields.get("chunks_done", 0) elif event == "file_done": eta["done_chars"] += eta["cur_chars"] eta["cur_chars"] = 0 eta["cur_total"] = 0 eta["cur_done"] = 0 try: yield ( f"event: {event}\n" f"data: {json.dumps({'name': fname, **fields})}\n\n" ) except _disconnect_exceptions() as e: log.warning("process_stream: disconnect on file event sid=%s err=%r", sid, e) cancel.set() if cancel_event: cancel_event.set() return elif kind == "result": _, zip_data, csv_str, stats = evt log.info("process_stream: result sid=%s, storing result", sid) store_result(sid, zip_data) if csv_str: store_csv(sid, csv_str) count = 0 with zipfile.ZipFile(io.BytesIO(zip_data)) as zf: count = len(zf.namelist()) log.info("process_stream: complete sid=%s count=%d stats=%r", sid, count, stats) try: yield ( f"event: complete\n" f"data: {json.dumps({'total': count, **stats})}\n\n" ) except _disconnect_exceptions() as e: log.warning("process_stream: disconnect on complete sid=%s err=%r", sid, e) return return elif kind == "cancelled": _, zip_data, csv_str, stats, meta = evt log.info("process_stream: cancelled sid=%s, storing partial result", sid) store_result(sid, zip_data) if csv_str: store_csv(sid, csv_str) try: yield ( f"event: cancelled\n" f"data: {json.dumps({'saved': meta.get('processed', 0), 'total': meta.get('total', 0), **stats})}\n\n" ) except _disconnect_exceptions() as e: log.warning("process_stream: disconnect on cancelled sid=%s err=%r", sid, e) return return elif kind == "error": _, error_data = evt log.error("process_stream: error event sid=%s llm_incomplete=%s", sid, error_data.get("llm_incomplete", False)) try: yield f"event: proc_error\ndata: {json.dumps(error_data)}\n\n" except _disconnect_exceptions() as e: log.warning("process_stream: disconnect on error sid=%s err=%r", sid, e) return return finally: resume_ttl(sid) log.debug("process_stream: generator exit sid=%s, TTL resumed", sid) return Response( stream_with_context(generate()), content_type="text/event-stream", headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"} ) @api_bp.route("/process/", methods=["POST"]) def process(sid): """Process all session files -> ZIP (единым вызовом, общий mapping).""" files = get_files(sid) if files is None: return jsonify({"ok": False, "error": "Session not found"}), 404 if not files: return jsonify({"ok": False, "error": "No files"}), 400 store_result(sid, None) store_csv(sid, None) try: llm = LLMClient() all_files = [(fname, content, "") for fname, content in files] pause_ttl(sid) try: zip_data, csv_str, _ = obfuscate_files(all_files, llm_client=llm) finally: resume_ttl(sid) store_result(sid, zip_data) if csv_str: store_csv(sid, csv_str) return jsonify({"ok": True, "status": "done"}) except LLMIncompleteError as e: log.exception("process: LLM incomplete sid=%s chunks=%d", sid, e.llm_errors) return jsonify({ "ok": False, "error": "Processing stopped because the LLM could not process all fragments", "llm_incomplete": True, "llm_errors": e.llm_errors, }), 500 except Exception: log.exception("process: processing failed sid=%s", sid) return jsonify({"ok": False, "error": "Processing failed"}), 500 @api_bp.route("/download/", methods=["GET"]) def download(sid): """Download result and cleanup session.""" zip_data = get_result(sid) if zip_data is None: return jsonify({"ok": False, "error": "Not found"}), 404 ts = (datetime.now() + timedelta(hours=3)).strftime("%Y-%m-%d_%H-%M-%S") return send_file(io.BytesIO(zip_data), mimetype="application/zip", as_attachment=True, download_name=f"obfuscator_{ts}.zip") @api_bp.route("/csv/", methods=["GET"]) def csv_download(sid): """Download CSV separately.""" csv_str = get_csv(sid) if csv_str is None: return jsonify({"ok": False, "error": "Not found"}), 404 ts = (datetime.now() + timedelta(hours=3)).strftime("%Y-%m-%d_%H-%M-%S") buf = io.BytesIO() buf.write('\ufeff'.encode('utf-8') + csv_str.encode('utf-8')) buf.seek(0) return send_file(buf, mimetype="text/csv", as_attachment=True, download_name=f"mapping_{ts}.csv")