Initial deployable Flask application

drhider/extractor.py

@@ -0,0 +1,357 @@
+"""
+Извлечение текста из документов — все форматы → Markdown.
+
+Поддерживает:
+- .docx → MD (python-docx: стили, жирный/курсив, таблицы)
+- .pdf → MD (pdfplumber: текст + таблицы, без форматирования)
+- .doc → MD (LibreOffice headless: .doc → .docx → штатный docx_to_markdown)
+- .txt и прочие → как есть
+
+Также содержит:
+- expand_zips — распаковка ZIP с защитой от ZIP-бомб
+"""
+
+import io
+import os
+import zipfile
+import tempfile
+import shutil
+import subprocess
+import logging
+from typing import List, Tuple
+
+log = logging.getLogger("drhider")
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# 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("")
+
+            # ── Таблицы ──
+            # Быстрый фильтр: если на странице нет линий/кривых/прямоугольников,
+            # таблиц нет (pdfplumber ищет таблицы только по векторным линиям).
+            # Включая page.rects — таблицы на цветном фоне (заливка).
+            # Пропускаем extract_tables() на «чистых» страницах (скан/текст) — это
+            # основная стоимость на сканах, где линий нет.
+            if page.lines or page.curves or page.rects:
+                tables = page.extract_tables()
+            else:
+                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)
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# DOC → Markdown (через сервис liberta)
+# ═══════════════════════════════════════════════════════════════════════════
+
+def doc_to_markdown(content: bytes) -> str:
+    """Конвертировать бинарный .doc → .docx (сервис liberta) → Markdown.
+
+    Отправляет .doc на HTTP-сервис конвертации, получает .docx,
+    затем прогоняет через штатный docx_to_markdown().
+
+    Args:
+        content: Бинарное содержимое .doc файла
+
+    Returns:
+        Строка в формате Markdown (или сообщение об ошибке)
+    """
+    import httpx
+    url = os.environ.get(
+        "CONVERT_SERVICE_URL",
+        "https://liberta.containerk8s.dev.nubes.ru/convert"
+    )
+    try:
+        with httpx.Client(timeout=120) as client:
+            resp = client.post(
+                url,
+                files={"file": ("input.doc", content, "application/msword")}
+            )
+        if resp.status_code == 200:
+            return docx_to_markdown(resp.content)
+        err = resp.json().get("error", f"HTTP {resp.status_code}")
+        log.warning("liberta conversion failed: %s", err)
+        return f"[DOC — conversion failed: {err}]"
+    except httpx.TimeoutException:
+        return "[DOC — conversion timed out (>120s)]"
+    except Exception as e:
+        log.warning("liberta error: %s", e)
+        return f"[DOC — conversion error: {e}]"
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# Единая точка входа
+# ═══════════════════════════════════════════════════════════════════════════
+
+def extract_text(fname: str, content: bytes, ctype: str = "") -> str:
+    """Извлечь текст из одного файла → Markdown.
+
+    Args:
+        fname: Имя файла (с расширением)
+        content: Бинарное содержимое файла
+        ctype: MIME-тип (зарезервировано, не используется)
+
+    Returns:
+        Строка в формате Markdown
+    """
+    ext = os.path.splitext(fname)[1].lower()
+
+    # ── .docx: полный Markdown с форматированием ──
+    if ext == '.docx':
+        try:
+            return docx_to_markdown(content)
+        except ImportError:
+            return content.decode('utf-8', errors='replace')
+
+    # ── .pdf: текст + таблицы ──
+    elif ext == '.pdf':
+        try:
+            return pdf_to_markdown(content)
+        except ImportError:
+            return content.decode('utf-8', errors='replace')
+
+    # ── .doc: бинарный → LibreOffice → DOCX → Markdown ──
+    elif ext == '.doc':
+        return doc_to_markdown(content)
+
+    # ── .txt и прочие ──
+    else:
+        return content.decode('utf-8', errors='replace')
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# ZIP-распаковка (без изменений)
+# ═══════════════════════════════════════════════════════════════════════════
+
+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 раскрыты, остальные как есть)
+    """
+    # Лимиты защиты от 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. Пытаемся восстановить исходные
+        байты (encode cp437) и декодировать:
+          1) как UTF-8 — если результат в кириллице/ASCII (Info-ZIP/Linux-зипы
+             пишут UTF-8 без флага; случайная коллизия CP866→UTF-8 отсекается
+             проверкой диапазона);
+          2) как CP866 — реальные 1С-выгрузки.
+        """
+        def _cyr_ok(s: str) -> bool:
+            # Все символы — ASCII или кириллица (U+0400–U+04FF)
+            return all(ord(c) < 128 or 0x0400 <= ord(c) <= 0x04FF for c in s)
+
+        if not (info.flag_bits & 0x800) and any(ord(c) > 127 for c in name):
+            try:
+                raw = name.encode("cp437")
+            except UnicodeEncodeError:
+                return name
+            # 1) UTF-8 (Info-ZIP/Linux без флага)
+            try:
+                dec = raw.decode("utf-8")
+                if _cyr_ok(dec):
+                    return dec
+            except UnicodeDecodeError:
+                pass
+            # 2) CP866 (1С)
+            try:
+                return raw.decode("cp866")
+            except UnicodeDecodeError:
+                pass
+        return name
+
+    while queue:
+        fname, content, ctype = queue.pop(0)
+
+        if not fname.lower().endswith('.zip'):
+            result.append((fname, content, ctype))
+            continue
+
+        try:
+            with zipfile.ZipFile(io.BytesIO(content)) as zf:
+                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
+
+                # Проверяем лимиты ДО чтения содержимого (защита от бомб)
+                bomb = False
+                for info in infos:
+                    if info.compress_size > 0:
+                        ratio = info.file_size / info.compress_size
+                        if ratio > MAX_RATIO:
+                            bomb = True
+                            break
+                    if total_uncompressed + info.file_size > MAX_UNCOMPRESSED:
+                        bomb = True
+                        break
+
+                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:
+                        continue
+                    # Синхронизация с фронтом (listZipFiles): из zip берём ТОЛЬКО документы
+                    low = name.lower()
+                    if not (low.endswith('.zip') or low.endswith(('.pdf', '.doc', '.docx', '.txt', '.md'))):
+                        continue
+
+                    inner_data = zf.read(info)
+                    total_uncompressed += len(inner_data)
+                    if total_uncompressed > MAX_UNCOMPRESSED:
+                        log.warning("ZIP uncompressed limit exceeded (mid-read): %s", fname)
+                        break
+
+                    # Вложенные ZIP добавляем в очередь на повторную распаковку
+                    queue.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