""" Blueprint: API DrHider. Five endpoints: - POST /api/upload — upload one file -> {session_id} - 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 zipfile import traceback import logging from datetime import datetime, timedelta from flask import Blueprint, request, send_file, jsonify, Response, stream_with_context from drhider import obfuscate_files, LLMClient from session import (create_session, add_file, get_files, store_result, get_result, store_csv, get_csv, cleanup, file_count) 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", "") if not sid: sid = create_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 added = 0 had_unnamed = False for f in uploaded: if not f.filename: had_unnamed = True continue data = f.read() log.info("upload: sid=%s file=%r size=%d", sid, f.filename, len(data)) if not add_file(sid, f.filename, 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("/process_stream/", methods=["GET"]) def process_stream(sid): """SSE: process all session files, streaming per-file progress. Все файлы обрабатываются ЕДИНЫМ вызовом obfuscate_files (общий 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 all_files = [(fname, content, "") for fname, content in files] log.info("process_stream: start sid=%s files=%d", sid, len(all_files)) def generate(): llm = LLMClient() q = queue.Queue() cancel = threading.Event() def progress(phase, idx, total_, name, elapsed): q.put(("progress", phase, idx, name, total_, elapsed)) def worker(): log.info("worker: start sid=%s files=%d", sid, len(all_files)) t0 = datetime.utcnow() try: zip_data, csv_str = obfuscate_files( all_files, llm_client=llm, progress_cb=progress ) stats = { "tokens": llm.tokens_total, "llm_sec": round(llm.llm_sec, 1), } dt = (datetime.utcnow() - t0).total_seconds() 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 Exception as e: log.error("worker: exception sid=%s: %r\n%s", sid, e, traceback.format_exc()) q.put(("error", repr(e))) threading.Thread(target=worker, daemon=True).start() log.debug("process_stream: worker thread started sid=%s", sid) 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 (для таймера в UI) try: yield ( f"event: llm\n" f"data: {json.dumps({'active': llm.llm_active, 'elapsed': round(llm.llm_elapsed_now(), 1), 'tokens': llm.tokens_total})}\n\n" ) except _disconnect_exceptions() as e: log.warning("process_stream: disconnect during heartbeat sid=%s err=%r", sid, e) cancel.set() return continue kind = evt[0] if kind == "progress": _, phase, idx, name, total_, elapsed = 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, '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() 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([n for n in zf.namelist() if n != "mapping.csv"]) 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 == "error": _, msg = evt log.error("process_stream: error event sid=%s msg=%r", sid, msg) try: yield f"event: error\ndata: {json.dumps({'error': msg})}\n\n" except _disconnect_exceptions() as e: log.warning("process_stream: disconnect on error sid=%s err=%r", sid, e) return return 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 try: llm = LLMClient() all_files = [(fname, content, "") for fname, content in files] zip_data, csv_str = obfuscate_files(all_files, llm_client=llm) store_result(sid, zip_data) if csv_str: store_csv(sid, csv_str) return jsonify({"ok": True, "status": "done"}) except Exception as e: traceback.print_exc() return jsonify({"ok": False, "error": str(e)}), 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"drhider_{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")