311 lines
12 KiB
Python
311 lines
12 KiB
Python
"""
|
|
Извлечение текста из документов — все форматы → 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("")
|
|
|
|
# ── Таблицы ──
|
|
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)
|
|
|
|
|
|
# ═══════════════════════════════════════════════════════════════════════════
|
|
# 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 раскрыты, остальные как есть)
|
|
"""
|
|
result: List[Tuple[str, bytes, str]] = []
|
|
|
|
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:
|
|
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
|
|
|
|
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))
|
|
break
|
|
|
|
name = info.filename
|
|
try:
|
|
name = name.encode("cp437").decode("utf-8", errors="replace")
|
|
except (UnicodeDecodeError, UnicodeEncodeError):
|
|
pass
|
|
|
|
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:
|
|
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
|