437 lines
18 KiB
Python
437 lines
18 KiB
Python
"""
|
||
Извлечение текста из документов и подготовка CSV к обфускации.
|
||
|
||
Поддерживает:
|
||
- .docx → MD (python-docx: стили, жирный/курсив, таблицы)
|
||
- .pdf → MD (pdfplumber: текст + таблицы, без форматирования)
|
||
- .doc → HTTP-сервис конвертации (CONVERT_SERVICE_URL) → .docx → Markdown
|
||
- .csv → текст ячеек для сканирования и обезличенный CSV на выходе
|
||
- .txt и прочие → как есть
|
||
|
||
Также содержит:
|
||
- expand_zips — распаковка ZIP с защитой от ZIP-бомб
|
||
"""
|
||
|
||
import io
|
||
import csv
|
||
import os
|
||
import zipfile
|
||
import tempfile
|
||
import shutil
|
||
import subprocess
|
||
import logging
|
||
from typing import Dict, List, Tuple
|
||
|
||
from . import replacer
|
||
|
||
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 decode_text(content: bytes) -> str:
|
||
"""Декодировать текст как UTF-8, затем CP1251, с Latin-1 последним fallback."""
|
||
try:
|
||
return content.decode("utf-8")
|
||
except UnicodeDecodeError:
|
||
# Блокнот и выгрузки 1С часто сохраняют кириллицу в CP1251.
|
||
try:
|
||
return content.decode("cp1251")
|
||
except UnicodeDecodeError:
|
||
return content.decode("latin-1", errors="replace")
|
||
|
||
|
||
def detect_delimiter(text: str) -> str:
|
||
"""Определить CSV-разделитель, используя Sniffer и детерминированный fallback."""
|
||
delimiters = [",", ";", "\t"]
|
||
nonempty_lines = [line for line in text.splitlines() if line.strip()]
|
||
if nonempty_lines:
|
||
sample = "\n".join(nonempty_lines[:10])
|
||
try:
|
||
return csv.Sniffer().sniff(sample, delimiters="".join(delimiters)).delimiter
|
||
except csv.Error:
|
||
first_line = nonempty_lines[0]
|
||
return max(delimiters, key=first_line.count)
|
||
return ","
|
||
|
||
|
||
def csv_to_text(content: bytes) -> str:
|
||
"""Распарсить CSV и склеить значения ячеек построчно для сканирования."""
|
||
text = decode_text(content).removeprefix("\ufeff")
|
||
delimiter = detect_delimiter(text)
|
||
reader = csv.reader(
|
||
io.StringIO(text, newline=""),
|
||
delimiter=delimiter,
|
||
quotechar='"',
|
||
strict=True,
|
||
)
|
||
# Каждая ячейка отделена переводом строки, чтобы сущности из соседних ячеек
|
||
# не объединялись в строку, по которой работают regex и LLM.
|
||
return "\n".join(cell for row in reader for cell in row)
|
||
|
||
|
||
def obfuscate_csv(content: bytes, mapping: Dict[str, str],
|
||
sorted_keys: List[str], compiled_re) -> bytes:
|
||
"""Заменить сущности в CSV-ячейках, сохранив порядок строк и столбцов."""
|
||
text = decode_text(content).removeprefix("\ufeff")
|
||
delimiter = detect_delimiter(text)
|
||
reader = csv.reader(
|
||
io.StringIO(text, newline=""),
|
||
delimiter=delimiter,
|
||
quotechar='"',
|
||
strict=True,
|
||
)
|
||
output = io.StringIO(newline="")
|
||
writer = csv.writer(
|
||
output,
|
||
delimiter=delimiter,
|
||
quotechar='"',
|
||
lineterminator="\n",
|
||
strict=True,
|
||
)
|
||
for row in reader:
|
||
writer.writerow([
|
||
replacer.apply_replacements(cell, mapping, sorted_keys, compiled_re)
|
||
for cell in row
|
||
])
|
||
|
||
# csv.writer нормализует кавычки/переводы строк; побайтовое сохранение
|
||
# исходного формата не является целью. BOM помогает Excel открыть кириллицу.
|
||
return ("\ufeff" + output.getvalue()).encode("utf-8")
|
||
|
||
|
||
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 decode_text(content)
|
||
|
||
# ── .pdf: текст + таблицы ──
|
||
elif ext == '.pdf':
|
||
try:
|
||
return pdf_to_markdown(content)
|
||
except ImportError:
|
||
return decode_text(content)
|
||
|
||
# ── .doc: бинарный → HTTP-сервис конвертации → DOCX → Markdown ──
|
||
elif ext == '.doc':
|
||
return doc_to_markdown(content)
|
||
|
||
# ── .csv: ячейки для обнаружения сущностей ──
|
||
elif ext == '.csv':
|
||
return csv_to_text(content)
|
||
|
||
# ── .txt и прочие ──
|
||
else:
|
||
return decode_text(content)
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════════
|
||
# 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', '.csv'))):
|
||
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
|