v0.0.32: фиксы по ревью — единый обфускатор, рекурсия ZIP, кириллица 1С, дедуп, SSE-дисконнект

drhider/builder.py

@@ -8,6 +8,7 @@
 
 import io
 import csv
+import os
 import zipfile
 import re
 import time
@@ -34,10 +35,30 @@ def build_zip(files: List[Tuple[str, bytes]], mapping_csv: str = "") -> bytes:
     """
     buf = io.BytesIO()
 
+    # Дедупликация имён: имя+одинаковый контент → пропуск; имя+разный контент → суффикс
+    seen: Dict[str, bytes] = {}
+
+    def _resolve_name(name: str, content: bytes) -> str:
+        if name not in seen:
+            seen[name] = content
+            return name
+        if seen[name] == content:
+            return None  # точный дубль — пропускаем
+        base, ext = os.path.splitext(name)
+        n = 2
+        while f"{base}_{n}{ext}" in seen:
+            n += 1
+        resolved = f"{base}_{n}{ext}"
+        seen[resolved] = content
+        return resolved
+
     with zipfile.ZipFile(buf, 'w', zipfile.ZIP_DEFLATED) as zf:
         # Добавляем обфусцированные файлы
         for fname, content in files:
-            info = zipfile.ZipInfo(fname)
+            out_name = _resolve_name(fname, content)
+            if out_name is None:
+                continue
+            info = zipfile.ZipInfo(out_name)
             info.date_time = time.localtime(time.time() + 3 * 3600)[:6]  # MSK (UTC+3)
             info.flag_bits |= 0x800  # Флаг: имя файла в UTF-8
             zf.writestr(info, content)

drhider/extractor.py

@@ -247,61 +247,73 @@ def expand_zips(files: List[Tuple[str, bytes, str]]) -> List[Tuple[str, bytes, s
     Returns:
         Новый список файлов (ZIP раскрыты, остальные как есть)
     """
-    result: List[Tuple[str, bytes, str]] = []
+    # Лимиты защиты от ZIP-бомб (глобально на весь вызов)
+    MAX_FILES_IN_ARCHIVE = 500
+    MAX_UNCOMPRESSED = 500 * 1024 * 1024  # 500 MB
+    MAX_RATIO = 100
+
+    result: List[Tuple[str, bytes, str]] = []
+    queue: List[Tuple[str, bytes, str]] = list(files)
+    total_uncompressed = 0
+
+    def _decode_name(name: str, info) -> str:
+        """Декодировать имя из ZIP.
+
+        Если флаг UTF-8 (bit 11) выставлен — имя уже корректное (берём как есть).
+        Иначе zipfile декодировал имя как CP437. Для кириллицы из 1С (CP866)
+        перекодируем CP437→CP866; при неудаче оставляем как есть.
+        """
+        if not (info.flag_bits & 0x800) and any(ord(c) > 127 for c in name):
+            try:
+                return name.encode("cp437").decode("cp866")
+            except (UnicodeDecodeError, UnicodeEncodeError):
+                pass
+        return name
+
+    while queue:
+        fname, content, ctype = queue.pop(0)
 
-    for fname, content, ctype in files:
         if not fname.lower().endswith('.zip'):
             result.append((fname, content, ctype))
             continue
 
         try:
             with zipfile.ZipFile(io.BytesIO(content)) as zf:
-                if len(zf.infolist()) > 500:
+                infos = [i for i in zf.infolist() if not i.is_dir()]
+
+                if len(infos) > MAX_FILES_IN_ARCHIVE:
                     log.warning("ZIP too many files, skipping: %s", fname)
                     result.append((fname, content, ctype))
                     continue
 
-                total_uncompressed = 0
-
-                for info in zf.infolist():
-                    if info.is_dir():
-                        continue
-
+                # Проверяем лимиты ДО чтения содержимого (защита от бомб)
+                bomb = False
+                for info in infos:
                     if info.compress_size > 0:
                         ratio = info.file_size / info.compress_size
-                        if ratio > 100:
-                            log.warning(
-                                "ZIP bomb ratio %.0f:1, skipping: %s", ratio, fname
-                            )
-                            result.append((fname, content, ctype))
+                        if ratio > MAX_RATIO:
+                            bomb = True
                             break
+                    if total_uncompressed + info.file_size > MAX_UNCOMPRESSED:
+                        bomb = True
+                        break
 
-                    name = info.filename
-                    try:
-                        name = name.encode("cp437").decode("utf-8", errors="replace")
-                    except (UnicodeDecodeError, UnicodeEncodeError):
-                        pass
+                if bomb:
+                    log.warning("ZIP bomb/limit, skipping: %s", fname)
+                    result.append((fname, content, ctype))
+                    continue
 
+                for info in infos:
+                    name = _decode_name(info.filename, info)
                     name = os.path.basename(name)
-                    if (
-                        not name
-                        or name.endswith("/")
-                        or ".." in name
-                        or "/" in name
-                        or "\\" in name
-                    ):
+                    if not name:
                         continue
 
                     inner_data = zf.read(info)
                     total_uncompressed += len(inner_data)
 
-                    if total_uncompressed > 500 * 1024 * 1024:
-                        log.warning(
-                            "ZIP uncompressed limit exceeded, stopping: %s", fname
-                        )
-                        break
-
-                    result.append((name, inner_data, ""))
+                    # Вложенные ZIP добавляем в очередь на повторную распаковку
+                    queue.append((name, inner_data, ""))
 
         except Exception as e:
             log.warning("Failed to expand ZIP %s: %s", fname, e)

drhider/obfuscator.py

@@ -20,6 +20,39 @@ from . import builder
 log = logging.getLogger("drhider")
 
 
+def _dedupe_file_names(
+    files: List[Tuple[str, bytes, str]]
+) -> List[Tuple[str, bytes, str]]:
+    """Уникализировать имена файлов перед обработкой.
+
+    Точный дубль (одинаковое имя + одинаковое содержимое) пропускается.
+    При коллизии имени с ДРУГИМ содержимым добавляется суффикс _2, _3, ...
+
+    Args:
+        files: [(filename, content_bytes, content_type), ...]
+
+    Returns:
+        Новый список с уникальными именами (порядок сохранён).
+    """
+    seen: Dict[str, bytes] = {}
+    out: List[Tuple[str, bytes, str]] = []
+
+    for fname, content, ctype in files:
+        candidate = fname
+        if candidate in seen:
+            if seen[candidate] == content:
+                continue  # точный дубль — пропускаем
+            base, ext = os.path.splitext(fname)
+            n = 2
+            while f"{base}_{n}{ext}" in seen:
+                n += 1
+            candidate = f"{base}_{n}{ext}"
+        seen[candidate] = content
+        out.append((candidate, content, ctype))
+
+    return out
+
+
 class TwoPassObfuscator:
     """Двухпроходный обфускатор документов.
 
@@ -52,7 +85,8 @@ class TwoPassObfuscator:
     # ═══════════════════════════════════════════════════════════════════
 
     def obfuscate(
-        self, files: List[Tuple[str, bytes, str]]
+        self, files: List[Tuple[str, bytes, str]],
+        progress_cb: Optional[Callable[[str, int, int, str], None]] = None
     ) -> Tuple[bytes, str]:
         """Обфусцировать список файлов.
 
@@ -60,14 +94,18 @@ class TwoPassObfuscator:
 
         Args:
             files: [(filename, content_bytes, content_type), ...]
+            progress_cb: Опциональный коллбек (phase, idx, total, fname),
+                         где phase ∈ {"start", "done"}, idx — 0-based индекс.
+                         Вызывается вокруг финальной обработки каждого файла.
 
         Returns:
             (zip_bytes, csv_string):
                 zip_bytes — ZIP-архив с обфусцированными .md файлами + mapping.csv
                 csv_string — содержимое mapping.csv как строка
         """
-        # ── Предобработка: распаковать ZIP ──
+        # ── Предобработка: распаковать ZIP + уникализировать имена ──
         files = extractor.expand_zips(files)
+        files = _dedupe_file_names(files)
 
         try:
             # ── Проход 1: сбор сущностей ──
@@ -93,8 +131,14 @@ class TwoPassObfuscator:
 
             # ── Проход 2: замена сущностей ──
             results: List[Tuple[str, bytes]] = []
+            total = len(files)
+
+            for i, (in_fname, content, ctype) in enumerate(files):
+                fname = in_fname
+                display_name = os.path.basename(in_fname) or in_fname
+                if progress_cb:
+                    progress_cb("start", i, total, display_name)
 
-            for fname, content, ctype in files:
                 obf_content = content  # По умолчанию — без изменений
 
                 if fname.endswith('.doc'):
@@ -118,6 +162,8 @@ class TwoPassObfuscator:
                     fname = md_name
 
                 results.append((fname, obf_content))
+                if progress_cb:
+                    progress_cb("done", i, total, display_name)
 
             # ── Сборка результата ──
             csv_str = builder.build_mapping_csv(self._mapping)
@@ -137,7 +183,8 @@ class TwoPassObfuscator:
 # ═══════════════════════════════════════════════════════════════════════
 
 def obfuscate_files(
-    files: List[Tuple[str, bytes, str]], llm_client=None
+    files: List[Tuple[str, bytes, str]], llm_client=None,
+    progress_cb: Optional[Callable[[str, int, int, str], None]] = None
 ) -> Tuple[bytes, str]:
     """Обфусцировать список файлов — удобная функция.
 
@@ -146,9 +193,10 @@ def obfuscate_files(
     Args:
         files: [(filename, content_bytes, content_type), ...]
         llm_client: Опциональный LLM-клиент
+        progress_cb: Опциональный коллбек (phase, idx, total, fname)
 
     Returns:
         (zip_bytes, csv_string)
     """
     obf = TwoPassObfuscator(llm_client=llm_client)
-    return obf.obfuscate(files)
+    return obf.obfuscate(files, progress_cb=progress_cb)

site/app.py

@@ -20,7 +20,7 @@ if _sys_path_root not in sys.path:
     sys.path.insert(0, _sys_path_root)
 
 # Версия приложения (меняется при изменениях)
-VERSION = "0.0.31"
+VERSION = "0.0.32"
 
 
 def create_app():

site/routes/api_bp.py

@@ -11,84 +11,130 @@ Five endpoints:
 
 import io
 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 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")
 
 
+def _disconnect_exceptions():
+    """Исключения, означающие отключение клиента SSE."""
+    return (GeneratorExit, BrokenPipeError, ConnectionResetError)
+
+
 @api_bp.route("/upload", methods=["POST"])
 def upload():
-    """Upload one file to session."""
+    """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
-    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
+
+    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_stream/<sid>", methods=["GET"])
 def process_stream(sid):
-    """SSE: process all session files one by one, streaming per-file progress."""
+    """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
 
-    total = len(files)
+    all_files = [(fname, content, "") for fname, content in 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  # клиент отключился
+        q = queue.Queue()
+        cancel = threading.Event()
 
-            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  # клиент отключился, НЕ обрабатываем остальные файлы
+        def progress(phase, idx, total_, name):
+            q.put(("progress", phase, idx, name, total_))
 
-        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"
+        def worker():
+            try:
+                zip_data, csv_str = obfuscate_files(
+                    all_files, llm_client=llm, progress_cb=progress
+                )
+                q.put(("result", zip_data, csv_str))
+            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
+                continue
+
+            kind = evt[0]
+
+            if kind == "progress":
+                _, phase, idx, name, total_ = evt
+                try:
+                    yield (
+                        f"event: {phase}\n"
+                        f"data: {json.dumps({'idx': idx, 'name': name, 'total': total_})}\n\n"
+                    )
+                except _disconnect_exceptions():
+                    cancel.set()
+                    return
+
+            elif kind == "result":
+                _, zip_data, csv_str = 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\ndata: {json.dumps({'total': count})}\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()),
@@ -99,7 +145,7 @@ def process_stream(sid):
 
 @api_bp.route("/process/<sid>", methods=["POST"])
 def process(sid):
-    """Process all session files one by one -> ZIP."""
+    """Process all session files -> ZIP (единым вызовом, общий mapping)."""
     files = get_files(sid)
     if files is None:
         return jsonify({"ok": False, "error": "Session not found"}), 404
@@ -107,28 +153,12 @@ def process(sid):
         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)
+        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", "files": len(all_results)})
+        return jsonify({"ok": True, "status": "done"})
     except Exception as e:
         traceback.print_exc()
         return jsonify({"ok": False, "error": str(e)}), 500

site/session.py

@@ -21,6 +21,9 @@ from typing import Dict, List, Tuple, Optional
 
 TTL_SECONDS = 30 * 60  # 30 минут
 
+# Максимальный суммарный объём файлов в одной сессии (защита памяти)
+MAX_SESSION_BYTES = 500 * 1024 * 1024  # 500 MB
+
 _sessions: Dict[str, dict] = {}
 _lock = threading.Lock()
 
@@ -73,6 +76,9 @@ def add_file(sid: str, filename: str, content: bytes) -> bool:
         s = _sessions.get(sid)
         if not s:
             return False
+        total = sum(len(c) for _, c in s["files"])
+        if total + len(content) > MAX_SESSION_BYTES:
+            return False  # превышен суммарный лимит сессии
         s["files"].append((filename, content))
         return True