Фаза 2: llm_client.py, extractor.py, scanner.py, replacer.py, builder.py, obfuscator.py, __init__.py

drhider/scanner.py

@@ -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)