Фаза 2: llm_client.py, extractor.py, scanner.py, replacer.py, builder.py, obfuscator.py, __init__.py

drhider/extractor.py

@@ -0,0 +1,260 @@
+"""
+Извлечение текста из документов разных форматов.
+
+Поддерживает:
+- .docx (через python-docx)
+- .pdf (через pdfplumber)
+- .doc (бинарный — не парсится)
+- .txt и прочие (как UTF-8)
+
+Также содержит:
+- _expand_zips — распаковка ZIP с защитой от ZIP-бомб
+- _convert_pdfs_to_docx — конвертация PDF → DOCX
+"""
+
+import io
+import os
+import zipfile
+import logging
+from typing import Dict, List, Tuple, Optional
+
+log = logging.getLogger("drhider")
+
+
+def extract_text(fname: str, content: bytes, ctype: str) -> Tuple[str, Optional[object]]:
+    """Извлечь текст из одного файла.
+
+    Args:
+        fname: Имя файла (с расширением)
+        content: Бинарное содержимое файла
+        ctype: MIME-тип (не используется в текущей версии)
+
+    Returns:
+        (text, docx_document_or_None):
+            text — извлечённый текст (строка)
+            doc — объект python-docx Document или None
+    """
+    doc = None
+    text = ""
+    ext = os.path.splitext(fname)[1].lower()
+
+    # ── .docx: извлекаем текст + сохраняем Document для замен с форматированием ──
+    if ext == '.docx':
+        try:
+            from docx import Document
+        except ImportError:
+            # Если python-docx не установлен — читаем как plain text
+            text = content.decode('utf-8', errors='replace')
+            return text, None
+
+        doc = Document(io.BytesIO(content))
+        # Собираем текст из параграфов
+        paragraphs = [p.text for p in doc.paragraphs]
+        # Добавляем текст из таблиц
+        for table in doc.tables:
+            for row in table.rows:
+                row_text = " | ".join(cell.text for cell in row.cells)
+                paragraphs.append(row_text)
+        text = "\n".join(paragraphs)
+
+    # ── .pdf: извлекаем текст через pdfplumber ──
+    elif ext == '.pdf':
+        try:
+            import pdfplumber
+        except ImportError:
+            text = content.decode('utf-8', errors='replace')
+            return text, None
+
+        parts = []
+        with pdfplumber.open(io.BytesIO(content)) as pdf:
+            for page in pdf.pages:
+                # Текст страницы
+                t = page.extract_text()
+                if t:
+                    parts.append(t)
+                # Текст из таблиц
+                for table in page.extract_tables():
+                    for row in table:
+                        row_str = " | ".join(str(c) if c else "" for c in row)
+                        parts.append(row_str)
+        text = "\n".join(parts)
+
+    # ── .doc: бинарный формат — без libreoffice не парсим ──
+    elif ext == '.doc':
+        text = "[DOC binary — not parsed]"
+        return text, None
+
+    # ── .txt и прочие: читаем как UTF-8 ──
+    else:
+        text = content.decode('utf-8', errors='replace')
+
+    return text, doc
+
+
+def expand_zips(files: List[Tuple[str, bytes, str]]) -> List[Tuple[str, bytes, str]]:
+    """Распаковать ZIP-файлы в списке, заменив их содержимым.
+
+    Не-ZIP файлы проходят без изменений.
+
+    Защита от ZIP-бомб:
+    - Максимум 500 файлов в архиве
+    - Ratio file_size/compress_size не более 100:1
+    - Накопительный размер распакованных данных не более 500 MB
+
+    Args:
+        files: [(filename, content_bytes, content_type), ...]
+
+    Returns:
+        Новый список файлов (ZIP раскрыты, остальные как есть)
+    """
+    result: List[Tuple[str, bytes, str]] = []
+
+    for fname, content, ctype in files:
+        # Пропускаем не-ZIP
+        if not fname.lower().endswith('.zip'):
+            result.append((fname, content, ctype))
+            continue
+
+        try:
+            with zipfile.ZipFile(io.BytesIO(content)) as zf:
+                # Проверка: не более 500 файлов в архиве
+                if len(zf.infolist()) > 500:
+                    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
+
+                    # Проверка на ZIP-бомбу: ratio
+                    if info.compress_size > 0:
+                        ratio = info.file_size / info.compress_size
+                        if ratio > 100:  # Файл сжимается более чем в 100 раз
+                            log.warning(
+                                "ZIP bomb ratio %.0f:1, skipping: %s", ratio, fname
+                            )
+                            result.append((fname, content, ctype))
+                            break  # Пропускаем весь архив
+
+                    # Декодируем имя файла: cp437 → utf-8
+                    name = info.filename
+                    try:
+                        name = name.encode("cp437").decode("utf-8", errors="replace")
+                    except (UnicodeDecodeError, UnicodeEncodeError):
+                        pass
+
+                    # Защита от path traversal
+                    name = os.path.basename(name)
+                    if (
+                        not name
+                        or name.endswith("/")
+                        or ".." in name
+                        or "/" in name
+                        or "\\" in name
+                    ):
+                        continue
+
+                    # Читаем и проверяем накопительный размер
+                    inner_data = zf.read(info)
+                    total_uncompressed += len(inner_data)
+
+                    if total_uncompressed > 500 * 1024 * 1024:  # 500 MB
+                        log.warning(
+                            "ZIP uncompressed limit exceeded, stopping: %s", fname
+                        )
+                        break
+
+                    result.append((name, inner_data, ""))
+
+        except Exception as e:
+            log.warning("Failed to expand ZIP %s: %s", fname, e)
+            result.append((fname, content, ctype))
+
+    return result
+
+
+def convert_pdfs_to_docx(
+    files: List[Tuple[str, bytes, str]],
+) -> List[Tuple[str, bytes, str]]:
+    """Конвертировать PDF-файлы в DOCX через pdfplumber.
+
+    Не-PDF файлы проходят без изменений.
+    При совпадении имён к имени добавляется суффикс '_из_pdf'.
+
+    Конвертация:
+    - Текст страницы → параграфы DOCX
+    - Таблицы → таблицы DOCX со стилем 'Table Grid'
+
+    Args:
+        files: [(filename, content_bytes, content_type), ...]
+
+    Returns:
+        Новый список файлов (PDF заменены на DOCX)
+    """
+    import pdfplumber
+    from docx import Document as DocxDocument
+
+    result: List[Tuple[str, bytes, str]] = []
+    existing_names = {f[0] for f in files}
+
+    for fname, content, ctype in files:
+        # Пропускаем не-PDF
+        if not fname.lower().endswith('.pdf'):
+            result.append((fname, content, ctype))
+            continue
+
+        try:
+            doc = DocxDocument()
+
+            with pdfplumber.open(io.BytesIO(content)) as pdf:
+                for page in pdf.pages:
+                    # ── Таблицы ──
+                    tables = page.extract_tables()
+                    for table in tables:
+                        if not table:
+                            continue
+
+                        # Чистим строки: убираем полностью пустые
+                        rows = [
+                            [str(c or "").strip() for c in (row or [])]
+                            for row in table
+                        ]
+                        rows = [r for r in rows if any(r)]
+
+                        if rows:
+                            t = doc.add_table(rows=len(rows), cols=len(rows[0]))
+                            t.style = 'Table Grid'
+                            for ri, row in enumerate(rows):
+                                for ci, cell_text in enumerate(row):
+                                    t.rows[ri].cells[ci].text = cell_text
+
+                    # ── Текст ──
+                    text = page.extract_text()
+                    if text:
+                        for line in text.split('\n'):
+                            line = line.strip()
+                            if line:
+                                doc.add_paragraph(line)
+
+            # Сохраняем DOCX в буфер
+            buf = io.BytesIO()
+            doc.save(buf)
+
+            # Формируем новое имя
+            new_name = fname[:-4] + '.docx'
+            if new_name in existing_names:
+                new_name = fname[:-4] + '_из_pdf.docx'
+            existing_names.add(new_name)
+
+            result.append((new_name, buf.getvalue(), ctype))
+
+        except Exception as e:
+            log.warning("PDF→DOCX error for %s: %s", fname, e)
+            # При ошибке — оставляем оригинальный PDF
+            result.append((fname, content, ctype))
+
+    return result