264 lines
9.7 KiB
Python
264 lines
9.7 KiB
Python
"""
|
|
Blueprint: API DrHider.
|
|
|
|
Five endpoints:
|
|
- POST /api/upload — upload one file -> {session_id}
|
|
- GET /api/process_stream/<sid> — SSE: process files, per-file progress
|
|
- POST /api/process/<sid> — process all session files -> {status:"done"} (legacy)
|
|
- GET /api/download/<sid> — download ZIP (with timestamp name)
|
|
- GET /api/csv/<sid> — download CSV separately
|
|
|
|
Тестовый (только при TEST_INPUT_DIR):
|
|
- POST /api/process_names — обфускация файлов по именам из TEST_INPUT_DIR
|
|
"""
|
|
|
|
import io
|
|
import os
|
|
import json
|
|
import queue
|
|
import threading
|
|
import zipfile
|
|
import traceback
|
|
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")
|
|
|
|
# Тестовый API «по именам» доступен ТОЛЬКО если задана валидная папка TEST_INPUT_DIR.
|
|
TEST_INPUT_DIR = os.environ.get("TEST_INPUT_DIR", "").strip()
|
|
_TEST_API_ENABLED = bool(TEST_INPUT_DIR) and os.path.isdir(TEST_INPUT_DIR)
|
|
|
|
|
|
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:
|
|
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
|
|
if not add_file(sid, f.filename, f.read()):
|
|
return jsonify({"ok": False, "error": "Session not found"}), 404
|
|
added += 1
|
|
|
|
if added == 0:
|
|
err = "No filename" if had_unnamed else "No file"
|
|
return jsonify({"ok": False, "error": err}), 400
|
|
|
|
return jsonify({"ok": True, "session": sid, "count": file_count(sid)})
|
|
|
|
|
|
@api_bp.route("/process_names", methods=["POST"])
|
|
def process_names():
|
|
"""Тестовый API: обфускация файлов ПО ИМЕНАМ из TEST_INPUT_DIR.
|
|
|
|
Удобно для тестирования curl'ом (не таскать файлы заново).
|
|
|
|
Доступен ТОЛЬКО если задан и валиден TEST_INPUT_DIR (env).
|
|
Иначе endpoint возвращает 404 — как будто API нет.
|
|
|
|
Запрос: {"files": ["имя1.pdf", "имя2.pdf"]} (только имена, без путей)
|
|
Ответ: ZIP с обфусцированными .md + mapping.csv.
|
|
"""
|
|
if not _TEST_API_ENABLED:
|
|
return jsonify({"ok": False, "error": "Not found"}), 404
|
|
|
|
data = request.get_json(silent=True) or {}
|
|
names = data.get("files") or []
|
|
if not isinstance(names, list) or not names:
|
|
return jsonify({"ok": False, "error": "files: список имён файлов"}), 400
|
|
|
|
all_files = []
|
|
for name in names:
|
|
# Только basename — защита от path traversal
|
|
base = os.path.basename(str(name))
|
|
if not base:
|
|
continue
|
|
path = os.path.join(TEST_INPUT_DIR, base)
|
|
if not os.path.isfile(path):
|
|
return jsonify({"ok": False, "error": f"файл не найден: {base}"}), 404
|
|
with open(path, "rb") as fh:
|
|
all_files.append((base, fh.read(), ""))
|
|
|
|
if not all_files:
|
|
return jsonify({"ok": False, "error": "нет файлов"}), 400
|
|
|
|
try:
|
|
llm = LLMClient()
|
|
zip_data, csv_str = obfuscate_files(all_files, llm_client=llm)
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
return jsonify({"ok": False, "error": str(e)}), 500
|
|
|
|
return send_file(
|
|
io.BytesIO(zip_data),
|
|
mimetype="application/zip",
|
|
as_attachment=True,
|
|
download_name="drhider_test.zip",
|
|
)
|
|
|
|
|
|
@api_bp.route("/process_stream/<sid>", 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]
|
|
|
|
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():
|
|
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),
|
|
}
|
|
q.put(("result", zip_data, csv_str, stats))
|
|
except Exception as e:
|
|
q.put(("error", repr(e)))
|
|
|
|
threading.Thread(target=worker, daemon=True).start()
|
|
|
|
while True:
|
|
try:
|
|
evt = q.get(timeout=1)
|
|
except queue.Empty:
|
|
if cancel.is_set():
|
|
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():
|
|
cancel.set()
|
|
return
|
|
continue
|
|
|
|
kind = evt[0]
|
|
|
|
if kind == "progress":
|
|
_, phase, idx, name, total_, elapsed = evt
|
|
try:
|
|
yield (
|
|
f"event: {phase}\n"
|
|
f"data: {json.dumps({'idx': idx, 'name': name, 'total': total_, 'elapsed': elapsed})}\n\n"
|
|
)
|
|
except _disconnect_exceptions():
|
|
cancel.set()
|
|
return
|
|
|
|
elif kind == "result":
|
|
_, zip_data, csv_str, stats = evt
|
|
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"])
|
|
try:
|
|
yield (
|
|
f"event: complete\n"
|
|
f"data: {json.dumps({'total': count, **stats})}\n\n"
|
|
)
|
|
except _disconnect_exceptions():
|
|
return
|
|
return
|
|
|
|
elif kind == "error":
|
|
_, msg = evt
|
|
try:
|
|
yield f"event: error\ndata: {json.dumps({'error': msg})}\n\n"
|
|
except _disconnect_exceptions():
|
|
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/<sid>", 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/<sid>", 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/<sid>", 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")
|