Фаза 2: llm_client.py, extractor.py, scanner.py, replacer.py, builder.py, obfuscator.py, __init__.py
Deploy drhider / validate (push) Waiting to run
Deploy drhider / validate (push) Waiting to run
This commit is contained in:
@@ -0,0 +1,13 @@
|
||||
"""
|
||||
DrHider — обфускация документов (двухпроходная, в памяти, без БД).
|
||||
|
||||
Проход 1: собрать все сущности из всех файлов → глобальный словарь замен.
|
||||
Проход 2: применить замены → собрать ZIP с обфусцированными файлами + mapping.csv.
|
||||
|
||||
Согласованность: одна и та же сущность во всех файлах → одно и то же фиктивное значение.
|
||||
"""
|
||||
|
||||
from .obfuscator import TwoPassObfuscator, obfuscate_files
|
||||
from .llm_client import LLMClient
|
||||
|
||||
__all__ = ["TwoPassObfuscator", "obfuscate_files", "LLMClient"]
|
||||
@@ -0,0 +1,96 @@
|
||||
"""
|
||||
Сборка результата обфускации.
|
||||
|
||||
Два метода:
|
||||
1. _build_zip — упаковка обфусцированных файлов в ZIP
|
||||
2. _build_mapping_csv — генерация mapping.csv с таблицей замен
|
||||
"""
|
||||
|
||||
import io
|
||||
import csv
|
||||
import zipfile
|
||||
import re
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
from .config import ENTITY_PATTERNS, COMPANY_PATTERN
|
||||
|
||||
|
||||
def build_zip(files: List[Tuple[str, bytes]], mapping_csv: str = "") -> bytes:
|
||||
"""Собрать ZIP-архив с обфусцированными файлами.
|
||||
|
||||
В архив добавляются:
|
||||
- Все обфусцированные файлы (с оригинальными именами)
|
||||
- mapping.csv — таблица соответствия оригинал→замена (если не пустая)
|
||||
|
||||
Имена файлов в архиве — UTF-8 (бит 11 в flag_bits).
|
||||
|
||||
Args:
|
||||
files: [(filename, content_bytes), ...]
|
||||
mapping_csv: Строка CSV с таблицей замен (опционально)
|
||||
|
||||
Returns:
|
||||
Бинарное содержимое ZIP-архива
|
||||
"""
|
||||
buf = io.BytesIO()
|
||||
|
||||
with zipfile.ZipFile(buf, 'w', zipfile.ZIP_DEFLATED) as zf:
|
||||
# Добавляем обфусцированные файлы
|
||||
for fname, content in files:
|
||||
info = zipfile.ZipInfo(fname)
|
||||
info.flag_bits |= 0x800 # Флаг: имя файла в UTF-8
|
||||
zf.writestr(info, content)
|
||||
|
||||
# Добавляем mapping.csv с BOM (для корректного открытия в Excel)
|
||||
if mapping_csv:
|
||||
zf.writestr(
|
||||
"mapping.csv",
|
||||
'\ufeff'.encode('utf-8') + mapping_csv.encode('utf-8'),
|
||||
)
|
||||
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def build_mapping_csv(mapping: Dict[str, str]) -> str:
|
||||
"""Собрать mapping.csv — таблицу соответствия оригинал → замена.
|
||||
|
||||
Колонки:
|
||||
- тип_данных: тип сущности (phone, email, inn_ul, company, ...)
|
||||
- оригинал: исходное значение из документа
|
||||
- замена: фиктивное значение
|
||||
|
||||
Тип определяется проверкой каждого значения через ENTITY_PATTERNS
|
||||
и COMPANY_PATTERN. Если ни один паттерн не подошёл — тип "text".
|
||||
|
||||
Args:
|
||||
mapping: Словарь {оригинал: замена}
|
||||
|
||||
Returns:
|
||||
Строка в формате CSV
|
||||
"""
|
||||
buf = io.StringIO()
|
||||
writer = csv.writer(buf)
|
||||
writer.writerow(["тип_данных", "оригинал", "замена"])
|
||||
|
||||
# Порядок проверки типов: inn_fl ДО inn_ul (12 цифр vs 10)
|
||||
type_order = [
|
||||
"phone", "email", "inn_fl", "inn_ul", "ogrn",
|
||||
"kpp", "bik", "rs", "ks", "passport",
|
||||
]
|
||||
|
||||
for original, replacement in sorted(mapping.items()):
|
||||
# Определяем тип сущности
|
||||
etype = "text" # По умолчанию
|
||||
|
||||
for t in type_order:
|
||||
pat = ENTITY_PATTERNS.get(t, "")
|
||||
if pat and re.match(pat, original, re.IGNORECASE):
|
||||
etype = t
|
||||
break
|
||||
|
||||
# Компании проверяем отдельно (COMPANY_PATTERN не в ENTITY_PATTERNS)
|
||||
if COMPANY_PATTERN.match(original):
|
||||
etype = "company"
|
||||
|
||||
writer.writerow([etype, original, replacement])
|
||||
|
||||
return buf.getvalue()
|
||||
@@ -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
|
||||
@@ -0,0 +1,64 @@
|
||||
"""
|
||||
LLM-клиент для DrHider.
|
||||
|
||||
Используется в TwoPassObfuscator._scan_llm_ner() для обнаружения
|
||||
имён, компаний, адресов и паспортных данных через LLM API.
|
||||
|
||||
Интерфейс:
|
||||
client = LLMClient()
|
||||
result = client.complete(prompt) # str
|
||||
"""
|
||||
|
||||
import os
|
||||
import httpx
|
||||
|
||||
|
||||
class LLMClient:
|
||||
"""Клиент для вызова LLM API (aillm.ru / OpenAI-совместимый).
|
||||
|
||||
Использует переменные окружения:
|
||||
LLM_API_KEY или LLM_KEY — ключ API
|
||||
LLM_URL — URL эндпоинта (по умолчанию https://api.aillm.ru/v1/chat/completions)
|
||||
LLM_MODEL — модель (по умолчанию gpt-oss-120b)
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Инициализировать клиент. httpx импортируется лениво."""
|
||||
self._httpx = httpx
|
||||
|
||||
def complete(self, prompt: str) -> str:
|
||||
"""Отправить промпт в LLM и вернуть текст ответа.
|
||||
|
||||
Args:
|
||||
prompt: Текст промпта (инструкция + данные для анализа)
|
||||
|
||||
Returns:
|
||||
Текстовый ответ модели (обычно JSON-строка)
|
||||
|
||||
Raises:
|
||||
httpx.HTTPError: при ошибке HTTP
|
||||
KeyError: при неожиданном формате ответа
|
||||
"""
|
||||
# Ключ API: сначала LLM_KEY, затем LLM_API_KEY (для совместимости)
|
||||
key = os.environ.get("LLM_KEY") or os.environ.get("LLM_API_KEY", "")
|
||||
|
||||
# URL и модель с значениями по умолчанию
|
||||
url = os.environ.get("LLM_URL", "https://api.aillm.ru/v1/chat/completions")
|
||||
model = os.environ.get("LLM_MODEL", "gpt-oss-120b")
|
||||
|
||||
r = self._httpx.post(
|
||||
url,
|
||||
json={
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"max_tokens": 8000,
|
||||
"temperature": 0.1,
|
||||
},
|
||||
headers={
|
||||
"Authorization": f"Bearer {key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
timeout=120,
|
||||
)
|
||||
r.raise_for_status()
|
||||
return r.json()["choices"][0]["message"]["content"]
|
||||
@@ -0,0 +1,148 @@
|
||||
"""
|
||||
Оркестратор двухпроходной обфускации — класс TwoPassObfuscator.
|
||||
|
||||
Процесс:
|
||||
1. Проход 1 (сбор): извлечь текст → regex-сканирование → LLM-сканирование
|
||||
2. Проход 2 (замена): применить mapping ко всем файлам
|
||||
3. Сборка: ZIP + mapping.csv
|
||||
"""
|
||||
|
||||
import io
|
||||
import logging
|
||||
from typing import Dict, List, Tuple, Callable, Optional
|
||||
|
||||
from . import extractor
|
||||
from . import scanner
|
||||
from . import replacer
|
||||
from . import builder
|
||||
|
||||
log = logging.getLogger("drhider")
|
||||
|
||||
|
||||
class TwoPassObfuscator:
|
||||
"""Двухпроходный обфускатор документов.
|
||||
|
||||
Проход 1: собрать все сущности из всех файлов → глобальный словарь замен.
|
||||
Проход 2: применить замены ко всем файлам → ZIP с результатом.
|
||||
|
||||
Согласованность: одна и та же сущность во всех файлах получает
|
||||
одно и то же фиктивное значение.
|
||||
|
||||
Attributes:
|
||||
_mapping: {оригинал: замена} — глобальный словарь
|
||||
_sorted_keys: ключи mapping, отсортированные по длине (убывание)
|
||||
_llm_client: опциональный LLM-клиент для NER
|
||||
"""
|
||||
|
||||
def __init__(self, llm_client=None):
|
||||
"""Инициализировать обфускатор.
|
||||
|
||||
Args:
|
||||
llm_client: Объект с методом .complete(prompt) -> str.
|
||||
Если None — LLM-сканирование не выполняется.
|
||||
"""
|
||||
self._mapping: Dict[str, str] = {}
|
||||
self._sorted_keys: List[str] = []
|
||||
self._llm_client = llm_client
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════
|
||||
# Главная точка входа
|
||||
# ═══════════════════════════════════════════════════════════════════
|
||||
|
||||
def obfuscate(
|
||||
self, files: List[Tuple[str, bytes, str]]
|
||||
) -> Tuple[bytes, str]:
|
||||
"""Обфусцировать список файлов.
|
||||
|
||||
Args:
|
||||
files: [(filename, content_bytes, content_type), ...]
|
||||
content_type — MIME-тип (может быть пустой строкой)
|
||||
|
||||
Returns:
|
||||
(zip_bytes, csv_string):
|
||||
zip_bytes — ZIP-архив с обфусцированными файлами + mapping.csv
|
||||
csv_string — содержимое mapping.csv как строка
|
||||
"""
|
||||
# ── Предобработка: распаковать ZIP, конвертировать PDF ──
|
||||
files = extractor.expand_zips(files)
|
||||
files = extractor.convert_pdfs_to_docx(files)
|
||||
|
||||
try:
|
||||
# ── Проход 1: сбор сущностей ──
|
||||
all_texts: Dict[str, str] = {}
|
||||
all_docx: Dict[str, object] = {}
|
||||
|
||||
for fname, content, ctype in files:
|
||||
text, doc = extractor.extract_text(fname, content, ctype)
|
||||
all_texts[fname] = text
|
||||
if doc is not None:
|
||||
all_docx[fname] = doc
|
||||
|
||||
# Regex-сканирование (быстрое, локальное)
|
||||
if text and text != "[DOC binary — not parsed]":
|
||||
scanner.scan_regex(text, self._mapping)
|
||||
|
||||
# LLM-сканирование (медленное, сетевое — только если есть клиент)
|
||||
if self._llm_client:
|
||||
scanner.scan_llm_ner(all_texts, self._mapping, self._llm_client)
|
||||
|
||||
# Предсортировать ключи один раз (по убыванию длины)
|
||||
self._sorted_keys = sorted(
|
||||
self._mapping.keys(), key=len, reverse=True
|
||||
)
|
||||
|
||||
# ── Проход 2: замена сущностей ──
|
||||
results: List[Tuple[str, bytes]] = []
|
||||
|
||||
for fname, content, ctype in files:
|
||||
obf_content = content # По умолчанию — без изменений
|
||||
|
||||
if fname.endswith('.doc'):
|
||||
# .doc — бинарный формат, оставляем как есть
|
||||
pass
|
||||
elif fname in all_docx:
|
||||
# DOCX: замена с сохранением форматирования
|
||||
obf_content = replacer.replace_in_docx(
|
||||
all_docx[fname], self._mapping, self._sorted_keys
|
||||
)
|
||||
else:
|
||||
# Plain text / PDF-текст: простая замена
|
||||
txt = all_texts.get(fname, '')
|
||||
obf_content = replacer.replace_in_text(
|
||||
txt, fname, self._mapping, self._sorted_keys
|
||||
)
|
||||
|
||||
results.append((fname, obf_content))
|
||||
|
||||
# ── Сборка результата ──
|
||||
csv_str = builder.build_mapping_csv(self._mapping)
|
||||
zip_data = builder.build_zip(results, csv_str)
|
||||
|
||||
return zip_data, csv_str
|
||||
|
||||
finally:
|
||||
# Очистка состояния (обфускатор может использоваться повторно)
|
||||
self._mapping.clear()
|
||||
self._sorted_keys.clear()
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# Удобная функция для быстрого вызова
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
def obfuscate_files(
|
||||
files: List[Tuple[str, bytes, str]], llm_client=None
|
||||
) -> Tuple[bytes, str]:
|
||||
"""Обфусцировать список файлов — удобная функция.
|
||||
|
||||
Создаёт экземпляр TwoPassObfuscator и вызывает .obfuscate().
|
||||
|
||||
Args:
|
||||
files: [(filename, content_bytes, content_type), ...]
|
||||
llm_client: Опциональный LLM-клиент
|
||||
|
||||
Returns:
|
||||
(zip_bytes, csv_string)
|
||||
"""
|
||||
obf = TwoPassObfuscator(llm_client=llm_client)
|
||||
return obf.obfuscate(files)
|
||||
@@ -0,0 +1,120 @@
|
||||
"""
|
||||
Проход 2: замена сущностей в документах.
|
||||
|
||||
Три метода:
|
||||
1. _apply_replacements — замена в plain-тексте (ядро)
|
||||
2. _replace_in_docx — замена в DOCX с сохранением форматирования
|
||||
3. _replace_in_text — замена в plain-тексте → bytes
|
||||
"""
|
||||
|
||||
import io
|
||||
import re
|
||||
from typing import Dict, List
|
||||
|
||||
|
||||
def apply_replacements(text: str, mapping: Dict[str, str], sorted_keys: List[str]) -> str:
|
||||
"""Применить все замены из словаря mapping к строке.
|
||||
|
||||
Ключи применяются в порядке убывания длины (sorted_keys).
|
||||
Это гарантирует, что более длинные совпадения заменяются раньше
|
||||
коротких (например, «ИНН 123456789012» до «ИНН 1234567890»).
|
||||
|
||||
Для сущностей, начинающихся и заканчивающихся на букву/цифру,
|
||||
используются границы слова (\\b), чтобы избежать частичных замен.
|
||||
|
||||
Args:
|
||||
text: Исходный текст
|
||||
mapping: Словарь {оригинал: замена}
|
||||
sorted_keys: Ключи mapping, отсортированные по длине (убывание)
|
||||
|
||||
Returns:
|
||||
Текст с заменами
|
||||
"""
|
||||
result = text
|
||||
|
||||
for original in sorted_keys:
|
||||
replacement = mapping[original]
|
||||
|
||||
# Если сущность обрамлена буквами/цифрами — используем границы слова
|
||||
if original and original[0].isalnum() and original[-1].isalnum():
|
||||
pattern = r'(?<!\w)' + re.escape(original) + r'(?!\w)'
|
||||
else:
|
||||
pattern = re.escape(original)
|
||||
|
||||
result = re.sub(pattern, replacement, result)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def replace_in_docx(doc, mapping: Dict[str, str], sorted_keys: List[str]) -> bytes:
|
||||
"""Заменить сущности в DOCX-документе с сохранением форматирования.
|
||||
|
||||
Алгоритм для каждого параграфа:
|
||||
1. Склеиваем текст всех runs в одну строку
|
||||
2. Применяем замены
|
||||
3. Пишем результат в первый run, очищаем остальные
|
||||
|
||||
Это нужно потому что python-docx разбивает текст на runs
|
||||
(например, mid-docx форматирование), и сущность может быть
|
||||
разорвана между несколькими runs.
|
||||
|
||||
Args:
|
||||
doc: Объект python-docx Document
|
||||
mapping: Словарь замен
|
||||
sorted_keys: Ключи mapping по убыванию длины
|
||||
|
||||
Returns:
|
||||
Бинарное содержимое изменённого DOCX-файла
|
||||
"""
|
||||
# ── Обработка параграфов ──
|
||||
for para in doc.paragraphs:
|
||||
if not para.runs:
|
||||
continue
|
||||
# Склеиваем все runs в одну строку
|
||||
full_text = "".join(run.text for run in para.runs)
|
||||
replaced = apply_replacements(full_text, mapping, sorted_keys)
|
||||
|
||||
# Если были замены — пишем в первый run, очищаем остальные
|
||||
if replaced != full_text:
|
||||
para.runs[0].text = replaced
|
||||
for run in para.runs[1:]:
|
||||
run.text = ""
|
||||
|
||||
# ── Обработка таблиц ──
|
||||
for table in doc.tables:
|
||||
for row in table.rows:
|
||||
for cell in row.cells:
|
||||
for para in cell.paragraphs:
|
||||
if not para.runs:
|
||||
continue
|
||||
full_text = "".join(run.text for run in para.runs)
|
||||
replaced = apply_replacements(full_text, mapping, sorted_keys)
|
||||
|
||||
if replaced != full_text:
|
||||
para.runs[0].text = replaced
|
||||
for run in para.runs[1:]:
|
||||
run.text = ""
|
||||
|
||||
# Сохраняем в буфер
|
||||
buf = io.BytesIO()
|
||||
doc.save(buf)
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def replace_in_text(text: str, fname: str, mapping: Dict[str, str], sorted_keys: List[str]) -> bytes:
|
||||
"""Заменить сущности в plain-тексте.
|
||||
|
||||
Для PDF и прочих нетекстовых форматов — отдаём текст.
|
||||
(MVP: без сохранения форматирования PDF).
|
||||
|
||||
Args:
|
||||
text: Исходный текст
|
||||
fname: Имя файла (для будущего использования — разные форматы)
|
||||
mapping: Словарь замен
|
||||
sorted_keys: Ключи mapping по убыванию длины
|
||||
|
||||
Returns:
|
||||
Бинарное содержимое с заменами (UTF-8)
|
||||
"""
|
||||
replaced = apply_replacements(text, mapping, sorted_keys)
|
||||
return replaced.encode('utf-8')
|
||||
@@ -0,0 +1,132 @@
|
||||
"""
|
||||
Проход 1: обнаружение сущностей в тексте документов.
|
||||
|
||||
Два метода сканирования:
|
||||
1. _scan_regex — быстрое regex-обнаружение (телефоны, email, ИНН, счета, компании, ФИО)
|
||||
2. _scan_llm_ner — LLM-обнаружение (имена, адреса, паспорта — то что regex не ловит)
|
||||
"""
|
||||
|
||||
import re
|
||||
import json
|
||||
import logging
|
||||
from typing import Dict
|
||||
|
||||
from .config import ENTITY_PATTERNS, COMPANY_PATTERN, PERSON_PATTERN
|
||||
from .generators import ENTITY_GENERATORS
|
||||
from .generators.company import generate_company
|
||||
from .generators.person import generate_person
|
||||
from .generators.address import generate_address
|
||||
from .generators.passport import generate_passport
|
||||
|
||||
log = logging.getLogger("drhider")
|
||||
|
||||
|
||||
def scan_regex(text: str, mapping: Dict[str, str]) -> None:
|
||||
"""Сканировать текст regex-паттернами, заполнить словарь замен.
|
||||
|
||||
Ищет в тексте:
|
||||
- Сущности из ENTITY_PATTERNS (телефоны, email, ИНН, ОГРН, КПП, БИК, счета, паспорта)
|
||||
- Названия компаний (кроме «НУБЕС» — это Исполнитель, не заменяем)
|
||||
- ФИО в формате «Фамилия И.О.»
|
||||
|
||||
Найденные значения добавляются в mapping: оригинал → фиктивное значение.
|
||||
Если сущность уже есть в mapping — не перезаписываем (согласованность).
|
||||
|
||||
Args:
|
||||
text: Текст документа для сканирования
|
||||
mapping: Словарь замен (мутабельный, пополняется)
|
||||
"""
|
||||
# ── Сущности по regex-паттернам (телефоны, email, реквизиты) ──
|
||||
for entity_type, pattern in ENTITY_PATTERNS.items():
|
||||
for match in re.finditer(pattern, text, re.IGNORECASE | re.MULTILINE):
|
||||
original = match.group(0).strip()
|
||||
if not original or original in mapping:
|
||||
continue
|
||||
|
||||
# Выбираем генератор по типу сущности
|
||||
generator = ENTITY_GENERATORS.get(entity_type, lambda x: "XXX")
|
||||
mapping[original] = generator(original)
|
||||
|
||||
# ── Названия компаний ──
|
||||
for match in COMPANY_PATTERN.finditer(text):
|
||||
original = match.group(0).strip()
|
||||
if not original or original in mapping:
|
||||
continue
|
||||
|
||||
# «НУБЕС» — Исполнитель, НЕ заменяем
|
||||
if re.search(r'НУБЕС|NUBES', original, re.IGNORECASE):
|
||||
continue
|
||||
|
||||
mapping[original] = generate_company(original)
|
||||
|
||||
# ── ФИО (Фамилия И.О.) ──
|
||||
for match in PERSON_PATTERN.finditer(text):
|
||||
original = match.group(0).strip()
|
||||
if not original or original in mapping:
|
||||
continue
|
||||
mapping[original] = generate_person(original)
|
||||
|
||||
|
||||
def scan_llm_ner(all_texts: Dict[str, str], mapping: Dict[str, str], llm_client) -> None:
|
||||
"""LLM NER для обнаружения имён, адресов, паспортных данных.
|
||||
|
||||
Отправляет объединённый текст всех файлов в LLM, получает JSON-список
|
||||
найденных сущностей, добавляет их в словарь замен.
|
||||
|
||||
Отправляются первые 3000 символов каждого файла (экономия токенов).
|
||||
Общий размер промпта ограничен 8000 символами.
|
||||
|
||||
Args:
|
||||
all_texts: {filename: text_content} — тексты всех файлов
|
||||
mapping: Словарь замен (мутабельный, пополняется)
|
||||
llm_client: Объект с методом .complete(prompt) -> str
|
||||
"""
|
||||
# Формируем комбинированный текст: имя файла + первые 3000 символов
|
||||
combined = "\n\n---FILE---\n\n".join(
|
||||
f"FILE: {fname}\n{t[:3000]}" for fname, t in all_texts.items()
|
||||
)
|
||||
|
||||
# Промпт для LLM
|
||||
prompt = (
|
||||
"Ты — система обнаружения персональных данных в документах. "
|
||||
"Найди ВСЕ следующие сущности в тексте ниже:\n\n"
|
||||
"1. ФИО (полные и сокращённые — 'Иванов И.И.', 'Петров А.С.')\n"
|
||||
"2. Названия компаний-контрагентов (не 'НУБЕС')\n"
|
||||
"3. Почтовые адреса\n"
|
||||
"4. Паспортные данные\n\n"
|
||||
"Формат ответа — JSON-массив:\n"
|
||||
'[{"type": "person"|"company"|"address"|"passport", "value": "найденный текст"}]\n\n'
|
||||
f"Текст:\n{combined[:8000]}"
|
||||
)
|
||||
|
||||
try:
|
||||
# Отправляем запрос в LLM
|
||||
raw = llm_client.complete(prompt)
|
||||
|
||||
# Чистим markdown-обёртку (если LLM вернула ```json ... ```)
|
||||
raw = raw.strip()
|
||||
if raw.startswith("```"):
|
||||
raw = raw.split("\n", 1)[1]
|
||||
if raw.endswith("```"):
|
||||
raw = raw[:-3]
|
||||
|
||||
entities = json.loads(raw)
|
||||
|
||||
# Добавляем найденные сущности в mapping
|
||||
for ent in entities:
|
||||
val = ent.get("value", "").strip()
|
||||
if not val or val in mapping:
|
||||
continue
|
||||
|
||||
ent_type = ent.get("type", "")
|
||||
if ent_type == "person":
|
||||
mapping[val] = generate_person(val)
|
||||
elif ent_type == "company":
|
||||
mapping[val] = generate_company(val)
|
||||
elif ent_type == "address":
|
||||
mapping[val] = generate_address(val)
|
||||
elif ent_type == "passport":
|
||||
mapping[val] = generate_passport(val)
|
||||
|
||||
except Exception as e:
|
||||
log.warning("LLM NER failed: %s", e)
|
||||
Reference in New Issue
Block a user