extractor.py → MD-пайплайн (docx_to_markdown, pdf_to_markdown), obfuscator.py — новый пайплайн без docx-объектов

drhider/extractor.py

@@ -1,95 +1,193 @@
 """
-Извлечение текста из документов разных форматов.
+Извлечение текста из документов — все форматы → Markdown.
 
 Поддерживает:
-- .docx (через python-docx)
-- .pdf (через pdfplumber)
-- .doc (бинарный — не парсится)
-- .txt и прочие (как UTF-8)
+- .docx → MD (python-docx: стили, жирный/курсив, таблицы)
+- .pdf → MD (pdfplumber: текст + таблицы, без форматирования)
+- .doc (бинарный — не парсится, возвращает заглушку)
+- .txt и прочие → как есть
 
 Также содержит:
-- _expand_zips — распаковка ZIP с защитой от ZIP-бомб
-- _convert_pdfs_to_docx — конвертация PDF → DOCX
+- expand_zips — распаковка ZIP с защитой от ZIP-бомб
 """
 
 import io
 import os
 import zipfile
 import logging
-from typing import Dict, List, Tuple, Optional
+from typing import List, Tuple
 
 log = logging.getLogger("drhider")
 
 
-def extract_text(fname: str, content: bytes, ctype: str) -> Tuple[str, Optional[object]]:
-    """Извлечь текст из одного файла.
+# ═══════════════════════════════════════════════════════════════════════════
+# DOCX → Markdown
+# ═══════════════════════════════════════════════════════════════════════════
+
+def docx_to_markdown(content: bytes) -> str:
+    """Конвертировать DOCX в Markdown с сохранением базового форматирования.
+
+    Маппинг стилей:
+    - Heading 1/2/3 → # / ## / ###
+    - Bold → **текст**
+    - Italic → *текст*
+    - Таблицы → | col1 | col2 | (с разделителем)
+
+    Args:
+        content: Бинарное содержимое .docx файла
+
+    Returns:
+        Строка в формате Markdown
+    """
+    from docx import Document
+
+    doc = Document(io.BytesIO(content))
+    lines = []
+
+    for para in doc.paragraphs:
+        if not para.text.strip():
+            lines.append("")  # Пустая строка — разделитель абзацев
+            continue
+
+        text = para.text
+
+        # Применяем форматирование из runs (bold, italic)
+        if para.runs:
+            formatted_parts = []
+            for run in para.runs:
+                part = run.text
+                if run.bold:
+                    part = f"**{part}**"
+                if run.italic:
+                    part = f"*{part}*"
+                formatted_parts.append(part)
+            text = "".join(formatted_parts)
+
+        # Стиль абзаца
+        style_name = para.style.name if para.style else ""
+
+        if style_name.startswith("Heading"):
+            # Heading 1 → #, Heading 2 → ##, ...
+            try:
+                level = int(style_name.split()[-1])
+            except (ValueError, IndexError):
+                level = 1
+            prefix = "#" * min(level, 6)
+            lines.append(f"{prefix} {text}")
+        else:
+            lines.append(text)
+
+    # ── Таблицы ──
+    for table in doc.tables:
+        lines.append("")  # Отступ перед таблицей
+        for ri, row in enumerate(table.rows):
+            cells = [cell.text.replace("\n", " ").strip() for cell in row.cells]
+            lines.append("| " + " | ".join(cells) + " |")
+            # Разделитель после первой строки (заголовка таблицы)
+            if ri == 0:
+                lines.append("|" + "|".join(["---"] * len(cells)) + "|")
+        lines.append("")  # Отступ после таблицы
+
+    return "\n".join(lines)
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# PDF → Markdown
+# ═══════════════════════════════════════════════════════════════════════════
+
+def pdf_to_markdown(content: bytes) -> str:
+    """Конвертировать PDF в Markdown (только текст + таблицы, без форматирования).
+
+    pdfplumber не различает стили (bold/italic/heading), поэтому
+    всё извлекается как plain text.
+
+    Args:
+        content: Бинарное содержимое .pdf файла
+
+    Returns:
+        Строка в формате Markdown (plain text + таблицы)
+    """
+    import pdfplumber
+
+    lines = []
+
+    with pdfplumber.open(io.BytesIO(content)) as pdf:
+        for page in pdf.pages:
+            # ── Текст страницы ──
+            text = page.extract_text()
+            if text:
+                lines.append(text)
+                lines.append("")
+
+            # ── Таблицы ──
+            tables = page.extract_tables()
+            for table in tables:
+                if not table:
+                    continue
+
+                # Чистим строки таблицы
+                clean_rows = [
+                    [str(c or "").strip() for c in (row or [])]
+                    for row in table
+                ]
+                clean_rows = [r for r in clean_rows if any(r)]
+
+                if not clean_rows:
+                    continue
+
+                # Выводим таблицу в MD
+                for ri, row in enumerate(clean_rows):
+                    lines.append("| " + " | ".join(row) + " |")
+                    if ri == 0:
+                        lines.append("|" + "|".join(["---"] * len(row)) + "|")
+                lines.append("")
+
+    return "\n".join(lines)
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# Единая точка входа
+# ═══════════════════════════════════════════════════════════════════════════
+
+def extract_text(fname: str, content: bytes, ctype: str = "") -> str:
+    """Извлечь текст из одного файла → Markdown.
 
     Args:
         fname: Имя файла (с расширением)
         content: Бинарное содержимое файла
-        ctype: MIME-тип (не используется в текущей версии)
+        ctype: MIME-тип (зарезервировано, не используется)
 
     Returns:
-        (text, docx_document_or_None):
-            text — извлечённый текст (строка)
-            doc — объект python-docx Document или None
+        Строка в формате Markdown
     """
-    doc = None
-    text = ""
     ext = os.path.splitext(fname)[1].lower()
 
-    # ── .docx: извлекаем текст + сохраняем Document для замен с форматированием ──
+    # ── .docx: полный Markdown с форматированием ──
     if ext == '.docx':
         try:
-            from docx import Document
+            return docx_to_markdown(content)
         except ImportError:
-            # Если python-docx не установлен — читаем как plain text
-            text = content.decode('utf-8', errors='replace')
-            return text, None
+            return content.decode('utf-8', errors='replace')
 
-        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 ──
+    # ── .pdf: текст + таблицы ──
     elif ext == '.pdf':
         try:
-            import pdfplumber
+            return pdf_to_markdown(content)
         except ImportError:
-            text = content.decode('utf-8', errors='replace')
-            return text, None
+            return content.decode('utf-8', errors='replace')
 
-        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 не парсим ──
+    # ── .doc: бинарный — пока не поддерживается ──
     elif ext == '.doc':
-        text = "[DOC binary — not parsed]"
-        return text, None
+        return "[DOC binary — not supported. Convert to .docx first.]"
 
-    # ── .txt и прочие: читаем как UTF-8 ──
+    # ── .txt и прочие ──
     else:
-        text = content.decode('utf-8', errors='replace')
+        return content.decode('utf-8', errors='replace')
 
-    return text, doc
 
+# ═══════════════════════════════════════════════════════════════════════════
+# ZIP-распаковка (без изменений)
+# ═══════════════════════════════════════════════════════════════════════════
 
 def expand_zips(files: List[Tuple[str, bytes, str]]) -> List[Tuple[str, bytes, str]]:
     """Распаковать ZIP-файлы в списке, заменив их содержимым.
@@ -110,14 +208,12 @@ def expand_zips(files: List[Tuple[str, bytes, str]]) -> List[Tuple[str, bytes, s
     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))
@@ -126,28 +222,24 @@ def expand_zips(files: List[Tuple[str, bytes, str]]) -> List[Tuple[str, bytes, s
                 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 раз
+                        if ratio > 100:
                             log.warning(
                                 "ZIP bomb ratio %.0f:1, skipping: %s", ratio, fname
                             )
                             result.append((fname, content, ctype))
-                            break  # Пропускаем весь архив
+                            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
@@ -158,11 +250,10 @@ def expand_zips(files: List[Tuple[str, bytes, str]]) -> List[Tuple[str, bytes, s
                     ):
                         continue
 
-                    # Читаем и проверяем накопительный размер
                     inner_data = zf.read(info)
                     total_uncompressed += len(inner_data)
 
-                    if total_uncompressed > 500 * 1024 * 1024:  # 500 MB
+                    if total_uncompressed > 500 * 1024 * 1024:
                         log.warning(
                             "ZIP uncompressed limit exceeded, stopping: %s", fname
                         )
@@ -175,86 +266,3 @@ def expand_zips(files: List[Tuple[str, bytes, str]]) -> List[Tuple[str, bytes, s
             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