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)