diff --git a/drhider/scanner.py b/drhider/scanner.py index 7dc007c..cd3c94d 100644 --- a/drhider/scanner.py +++ b/drhider/scanner.py @@ -2,8 +2,11 @@ Проход 1: обнаружение сущностей в тексте документов. Два метода сканирования: -1. _scan_regex — быстрое regex-обнаружение (телефоны, email, ИНН, счета, компании, ФИО) -2. _scan_llm_ner — LLM-обнаружение (имена, адреса, паспорта — то что regex не ловит) +1. scan_regex — быстрое regex-обнаружение (телефоны, email, ИНН, счета, компании, ФИО) +2. scan_llm_ner — универсальное LLM-обнаружение (любые приватные данные) + +LLM САМА определяет какие типы сущностей присутствуют в документе. +Никакого фиксированного списка типов. """ import re @@ -12,15 +15,25 @@ import logging from typing import Dict from .config import ENTITY_PATTERNS, COMPANY_PATTERN, PERSON_PATTERN -from .generators import ENTITY_GENERATORS +from .generators import ENTITY_GENERATORS, FALLBACK_GENERATOR 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 _get_generator(entity_type: str): + """Получить генератор по типу сущности. Если тип неизвестен — fallback. + + Args: + entity_type: Строковый идентификатор типа (например, "phone", "person_name") + + Returns: + Функция-генератор: (str) -> str + """ + return ENTITY_GENERATORS.get(entity_type, FALLBACK_GENERATOR) + + def scan_regex(text: str, mapping: Dict[str, str]) -> None: """Сканировать текст regex-паттернами, заполнить словарь замен. @@ -43,8 +56,8 @@ def scan_regex(text: str, mapping: Dict[str, str]) -> None: if not original or original in mapping: continue - # Выбираем генератор по типу сущности - generator = ENTITY_GENERATORS.get(entity_type, lambda x: "XXX") + # Генератор: специализированный (если есть) или fallback + generator = _get_generator(entity_type) mapping[original] = generator(original) # ── Названия компаний ── @@ -86,17 +99,21 @@ def scan_llm_ner(all_texts: Dict[str, str], mapping: Dict[str, str], llm_client) 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]}" + "You are a PII (Personally Identifiable Information) detector. " + "Find ALL sensitive or private information in the text below.\n\n" + "Return a JSON array. Each item must have:\n" + ' - "type": short snake_case label describing the entity\n' + ' (e.g. person_name, phone, email, address, inn, company, passport,\n' + ' contract_number, employee_id, bank_account — WHATEVER you see)\n' + ' - "value": the EXACT string from the text (copy verbatim)\n\n' + "Rules:\n" + "- Copy value VERBATIM — same spaces, punctuation, case as in text\n" + "- If same value appears multiple times — include only once\n" + "- Invent the type label yourself based on what you see\n" + "- Return ONLY the JSON array, no other text, no markdown wrapping\n\n" + f"Text:\n{combined[:8000]}" ) try: @@ -118,15 +135,12 @@ def scan_llm_ner(all_texts: Dict[str, str], mapping: Dict[str, str], llm_client) 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) + # Тип — любой (LLM сама придумала) + ent_type = ent.get("type", "").strip().lower().replace(" ", "_") + + # Генератор: специализированный (если тип совпал) или fallback + generator = _get_generator(ent_type) + mapping[val] = generator(val) except Exception as e: log.warning("LLM NER failed: %s", e)