""" 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 import httpx 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, MAX_FILE_BYTES) 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 len(data) > MAX_FILE_BYTES: log.warning("upload: file exceeds %dMB, skipped sid=%s file=%r size=%d", MAX_FILE_BYTES // (1024 * 1024), sid, f.filename, len(data)) continue 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("/upload_refs", methods=["POST"]) def upload_refs(): """Принять ссылки на файлы (загружены на ВМ-буфер), забрать по egress. Вход: JSON {"session": "...", "files": [{"name": str, "size": int, "url": str}]}. Каждый файл тянется ИСХОДЯЩИМ GET'ом с ВМ (egress не ограничен шлюзом), читается по частям (stream), кладётся в сессию. После успешного pull файл удаляется с ВМ (best-effort; TTL-чистка на ВМ тоже есть). """ data = request.get_json(silent=True) or {} sid = data.get("session") or create_session() refs = data.get("files") or [] if not refs: log.warning("upload_refs: no files, sid=%s", sid) return jsonify({"ok": False, "error": "No files"}), 400 added = 0 try: with httpx.Client(timeout=120, follow_redirects=True) as client: for ref in refs: name = ref.get("name") url = ref.get("url") if not name or not url: continue # Лимит на один файл (50 МБ): сверх лимита — пропускаем (не участвует) if (ref.get("size") or 0) > MAX_FILE_BYTES: log.warning("upload_refs: file exceeds %dMB, skip sid=%s file=%r size=%s", MAX_FILE_BYTES // (1024 * 1024), sid, name, ref.get("size")) try: client.delete(url) except Exception: pass continue with client.stream("GET", url) as resp: resp.raise_for_status() content = b"".join(resp.iter_bytes()) log.info("upload_refs: pulled sid=%s file=%r size=%d", sid, name, len(content)) if len(content) > MAX_FILE_BYTES: log.warning("upload_refs: pulled file exceeds %dMB, skip sid=%s file=%r size=%d", MAX_FILE_BYTES // (1024 * 1024), sid, name, len(content)) try: client.delete(url) except Exception: pass continue if not add_file(sid, name, content): log.warning("upload_refs: session not found/limit, sid=%s file=%r", sid, name) return jsonify({"ok": False, "error": "Session not found"}), 404 try: client.delete(url) # убрать файл с ВМ после загрузки except Exception: pass added += 1 except Exception as e: log.error("upload_refs: pull error sid=%s: %r", sid, e) return jsonify({"ok": False, "error": "Pull failed: %s" % e}), 502 log.info("upload_refs: 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): """Отладка: список файлов сессии с размерами (для теста загрузки).""" 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("/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")