drhider/site/routes/api_bp.py

"""
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
"""

import io
import json
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 drhider.builder import build_zip, build_mapping_csv
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_bp.route("/upload", methods=["POST"])
def upload():
    """Upload one file 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
    f = uploaded[0]
    if not f.filename:
        return jsonify({"ok": False, "error": "No filename"}), 400
    ok = add_file(sid, f.filename, f.read())
    if not ok:
        return jsonify({"ok": False, "error": "Session not found"}), 404
    return jsonify({"ok": True, "session": sid, "count": file_count(sid)})


@api_bp.route("/process_stream/<sid>", methods=["GET"])
def process_stream(sid):
    """SSE: process all session files one by one, streaming per-file progress."""
    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

    total = len(files)

    def generate():
        llm = LLMClient()
        all_results = []
        all_mapping = {}
        for idx, (fname, content) in enumerate(files):
            # Отправляем: начали файл
            try:
                yield f"event: start\ndata: {json.dumps({'idx': idx, 'name': fname, 'total': total})}\n\n"
            except GeneratorExit:
                return  # клиент отключился

            zip_data, _csv = obfuscate_files([(fname, content, "")], llm_client=llm)
            with zipfile.ZipFile(io.BytesIO(zip_data)) as zf:
                for name in zf.namelist():
                    if name == "mapping.csv":
                        csv_text = zf.read(name).decode("utf-8")
                        lines = csv_text.strip().split("\n")
                        if len(lines) > 1:
                            for line in lines[1:]:
                                parts = line.split(",", 2)
                                if len(parts) >= 3:
                                    all_mapping[parts[1]] = parts[2]
                    elif not name.endswith("/"):
                        all_results.append((name, zf.read(name)))
            # Отправляем: закончили файл (если клиент ушёл — освобождаем поток)
            try:
                yield f"event: done\ndata: {json.dumps({'idx': idx, 'name': fname, 'total': total})}\n\n"
            except GeneratorExit:
                return  # клиент отключился, НЕ обрабатываем остальные файлы

        csv_str = build_mapping_csv(all_mapping) if all_mapping else ""
        final_zip = build_zip(all_results)
        store_result(sid, final_zip)
        if csv_str:
            store_csv(sid, csv_str)
        yield f"event: complete\ndata: {json.dumps({'total': len(all_results)})}\n\n"

    return Response(
        stream_with_context(generate()),
        content_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no", "Connection": "close"}
    )


@api_bp.route("/process/<sid>", methods=["POST"])
def process(sid):
    """Process all session files one by one -> ZIP."""
    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_results = []
        all_mapping = {}
        for fname, content in files:
            zip_data, _csv = obfuscate_files([(fname, content, "")], llm_client=llm)
            with zipfile.ZipFile(io.BytesIO(zip_data)) as zf:
                for name in zf.namelist():
                    if name == "mapping.csv":
                        csv_text = zf.read(name).decode("utf-8")
                        lines = csv_text.strip().split("\n")
                        if len(lines) > 1:
                            for line in lines[1:]:
                                parts = line.split(",", 2)
                                if len(parts) >= 3:
                                    all_mapping[parts[1]] = parts[2]
                    elif not name.endswith("/"):
                        all_results.append((name, zf.read(name)))
        csv_str = build_mapping_csv(all_mapping) if all_mapping else ""
        final_zip = build_zip(all_results)  # CSV — отдельно, не в ZIP
        store_result(sid, final_zip)
        if csv_str:
            store_csv(sid, csv_str)
        return jsonify({"ok": True, "status": "done", "files": len(all_results)})
    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")