463 lines
20 KiB
Python
463 lines
20 KiB
Python
"""
|
||
DrHider — обфускация документов (двухпроходная, в памяти, без БД).
|
||
|
||
Проход 1: собрать все сущности из всех файлов → глобальный словарь замен.
|
||
Проход 2: применить замены → собрать ZIP с обфусцированными файлами + mapping.csv.
|
||
|
||
Согласованность: одна и та же сущность во всех файлах → одно и то же фиктивное значение.
|
||
"""
|
||
import io
|
||
import csv
|
||
import re
|
||
import hashlib
|
||
import random
|
||
import string
|
||
import zipfile
|
||
import logging
|
||
import tempfile
|
||
import os
|
||
import shutil
|
||
import datetime
|
||
from typing import Dict, List, Tuple, Callable, Optional
|
||
|
||
log = logging.getLogger("drhider")
|
||
|
||
# ═══════════════════════════════════════════
|
||
# Regex-паттерны для обнаружения сущностей
|
||
# ═══════════════════════════════════════════
|
||
|
||
ENTITY_PATTERNS: Dict[str, str] = {
|
||
"phone": r'(?:\+7|8)[\s\-]?\(?\d{3}\)?[\s\-]?\d{3}[\s\-]?\d{2}[\s\-]?\d{2}',
|
||
"email": r'[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}',
|
||
"inn_ul": r'ИНН\s*\d{10}',
|
||
"inn_fl": r'ИНН\s*\d{12}',
|
||
"ogrn": r'ОГРН\s*\d{13}',
|
||
"kpp": r'КПП\s*\d{9}',
|
||
"bik": r'БИК\s*\d{9}',
|
||
"rs": r'(?:р/с|расч[её]тный\s*сч[её]т)\s*\d{20}',
|
||
"ks": r'(?:к/с|корр?[еи]?спондентский\s*сч[её]т)\s*\d{20}',
|
||
"passport": r'\d{2}\s*\d{2}\s*\d{6,7}',
|
||
}
|
||
|
||
# Фирмы — обнаружение по шаблону
|
||
COMPANY_PATTERN = re.compile(
|
||
r'(?:ООО|ЗАО|ОАО|АО|ПАО|ИП|ТОО)\s+(?:«[^»]+»|"[^"]+"|[А-ЯA-Z][\w\s\-\.]+)',
|
||
re.IGNORECASE
|
||
)
|
||
|
||
# ═══════════════════════════════════════════
|
||
# Генераторы фиктивных значений
|
||
# ═══════════════════════════════════════════
|
||
|
||
# Словари русских имён
|
||
RU_SURNAMES = ["Иванов", "Смирнов", "Кузнецов", "Попов", "Васильев", "Петров",
|
||
"Соколов", "Михайлов", "Новиков", "Фёдоров", "Морозов", "Волков", "Алексеев",
|
||
"Лебедев", "Семёнов", "Егоров", "Павлов", "Козлов", "Степанов", "Николаев"]
|
||
RU_NAMES = ["Александр", "Дмитрий", "Сергей", "Андрей", "Алексей", "Максим",
|
||
"Евгений", "Иван", "Михаил", "Николай", "Владимир", "Павел", "Виктор", "Олег"]
|
||
RU_PATRONYMICS = ["Александрович", "Дмитриевич", "Сергеевич", "Андреевич",
|
||
"Алексеевич", "Иванович", "Михайлович", "Николаевич", "Владимирович",
|
||
"Павлович", "Викторович", "Олегович", "Евгеньевич", "Максимович"]
|
||
|
||
RU_CITIES = ["Москва", "Санкт-Петербург", "Новосибирск", "Екатеринбург",
|
||
"Казань", "Нижний Новгород", "Челябинск", "Самара", "Омск", "Ростов-на-Дону",
|
||
"Уфа", "Красноярск", "Воронеж", "Пермь", "Волгоград"]
|
||
RU_STREETS = ["Ленина", "Мира", "Пушкина", "Гагарина", "Советская",
|
||
"Кирова", "Октябрьская", "Молодёжная", "Садовая", "Центральная"]
|
||
|
||
FAKE_DOMAINS = ["example.ru", "mail.test", "company.local", "org.example.ru"]
|
||
|
||
|
||
def _checksum_inn10(inn: str) -> str:
|
||
"""Контрольная сумма для 10-значного ИНН."""
|
||
coeffs = [2, 4, 10, 3, 5, 9, 4, 6, 8]
|
||
s = sum(int(inn[i]) * coeffs[i] for i in range(9))
|
||
return str((s % 11) % 10)
|
||
|
||
|
||
def _checksum_inn12(inn: str) -> str:
|
||
"""Контрольные суммы для 12-значного ИНН."""
|
||
c1 = [7, 2, 4, 10, 3, 5, 9, 4, 6, 8]
|
||
c2 = [3, 7, 2, 4, 10, 3, 5, 9, 4, 6, 8]
|
||
s1 = sum(int(inn[i]) * c1[i] for i in range(10))
|
||
s2 = sum(int(inn[i]) * c2[i] for i in range(11))
|
||
return str((s1 % 11) % 10) + str((s2 % 11) % 10)
|
||
|
||
|
||
def _checksum_ogrn(ogrn: str) -> str:
|
||
"""Контрольная сумма для ОГРН (12 цифр → остаток от деления на 11)."""
|
||
s = int(ogrn) % 11
|
||
return str(s % 10)
|
||
|
||
|
||
def _random_digits(n: int) -> str:
|
||
return ''.join(random.choice(string.digits) for _ in range(n))
|
||
|
||
|
||
def _random_letters(n: int) -> str:
|
||
return ''.join(random.choice(string.ascii_lowercase) for _ in range(n))
|
||
|
||
|
||
def generate_phone(_: str) -> str:
|
||
code = random.choice(["495", "499", "812", "383", "343"])
|
||
return f"+7 ({code}) {_random_digits(3)}-{_random_digits(2)}-{_random_digits(2)}"
|
||
|
||
|
||
def generate_email(original: str) -> str:
|
||
"""Генерирует email с тем же форматом."""
|
||
domain = random.choice(FAKE_DOMAINS)
|
||
local = _random_letters(random.randint(5, 10))
|
||
return f"{local}@{domain}"
|
||
|
||
|
||
def generate_inn10(_: str) -> str:
|
||
base = f"{random.randint(1,9)}{_random_digits(8)}"
|
||
return base + _checksum_inn10(base)
|
||
|
||
|
||
def generate_inn12(_: str) -> str:
|
||
base = f"{random.randint(1,9)}{_random_digits(9)}"
|
||
return base + _checksum_inn12(base)
|
||
|
||
|
||
def generate_ogrn(_: str) -> str:
|
||
base = "1" + _random_digits(11)
|
||
return base + _checksum_ogrn(base)
|
||
|
||
|
||
def generate_kpp(_: str) -> str:
|
||
return _random_digits(4) + random.choice(["01", "43", "77"]) + _random_digits(3)
|
||
|
||
|
||
def generate_bik(_: str) -> str:
|
||
return "04" + _random_digits(7)
|
||
|
||
|
||
def generate_rs(_: str) -> str:
|
||
return "40702" + _random_digits(15)
|
||
|
||
|
||
def generate_ks(_: str) -> str:
|
||
return "30101" + _random_digits(15)
|
||
|
||
|
||
def generate_passport(_: str) -> str:
|
||
return f"{random.randint(10,99)} {random.randint(10,99)} {_random_digits(6)}"
|
||
|
||
|
||
def generate_company(_: str) -> str:
|
||
forms = ["ООО", "ЗАО", "АО"]
|
||
nouns = ["Технология", "Прогресс", "Гарант", "Стандарт", "Импульс",
|
||
"Вектор", "Сфера", "Альянс", "Синтез", "Меридиан", "Спектр", "Формат"]
|
||
return f'{random.choice(forms)} «{random.choice(nouns)}»'
|
||
|
||
|
||
def generate_person(_: str) -> str:
|
||
s = random.choice(RU_SURNAMES)
|
||
n = random.choice(RU_NAMES)
|
||
p = random.choice(RU_PATRONYMICS)
|
||
return f"{s} {n[0]}.{p[0]}."
|
||
|
||
|
||
def generate_address(_: str) -> str:
|
||
city = random.choice(RU_CITIES)
|
||
street = random.choice(RU_STREETS)
|
||
house = random.randint(1, 200)
|
||
return f"{city}, ул. {street}, д. {house}"
|
||
|
||
|
||
# ═══════════════════════════════════════════
|
||
# Основной класс
|
||
# ═══════════════════════════════════════════
|
||
|
||
class TwoPassObfuscator:
|
||
"""Двухпроходный обфускатор: сбор сущностей → замена."""
|
||
|
||
def __init__(self, llm_client=None):
|
||
self._mapping: Dict[str, str] = {} # оригинал → замена (только для точных текстовых совпадений)
|
||
self._regex_replacements: List[Tuple[str, str, Callable]] = [] # (pattern, type, generator)
|
||
self._llm_client = llm_client
|
||
self._workdir: Optional[str] = None
|
||
|
||
def obfuscate(self, files: List[Tuple[str, bytes, str]]) -> Tuple[bytes, str]:
|
||
"""
|
||
Главная точка входа.
|
||
|
||
Args:
|
||
files: [(original_filename, content_bytes, content_type), ...]
|
||
|
||
Returns:
|
||
(zip_bytes, mapping_csv_string)
|
||
"""
|
||
self._workdir = tempfile.mkdtemp(prefix="drhider_")
|
||
try:
|
||
# --- Распаковать ZIP-файлы ---
|
||
files = self._expand_zips(files)
|
||
|
||
# --- Проход 1: сбор сущностей ---
|
||
all_texts: Dict[str, str] = {} # filename → text for LLM NER
|
||
all_docx: Dict[str, object] = {} # filename → docx Document for replacement
|
||
|
||
for fname, content, ctype in files:
|
||
text, doc = self._extract_text(fname, content, ctype)
|
||
all_texts[fname] = text
|
||
if doc is not None:
|
||
all_docx[fname] = doc
|
||
self._scan_regex(text)
|
||
|
||
# LLM NER для имён/адресов (если есть клиент)
|
||
if self._llm_client:
|
||
self._scan_llm_ner(all_texts)
|
||
|
||
# --- Проход 2: замена ---
|
||
results = []
|
||
ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
seq = 0
|
||
for fname, content, ctype in files:
|
||
seq += 1
|
||
ext = os.path.splitext(fname)[1] or ".tmp"
|
||
clean_name = f"doc_{ts}_{seq:03d}{ext}"
|
||
obf_content = content
|
||
if fname in all_docx:
|
||
# Замена внутри docx
|
||
obf_content = self._replace_in_docx(all_docx[fname])
|
||
elif fname.endswith('.pdf'):
|
||
# PDF: read-only, создаём новый
|
||
obf_content = self._replace_in_text(all_texts.get(fname, ''), fname)
|
||
else:
|
||
# .doc или другой — замена по тексту
|
||
txt = all_texts.get(fname, '')
|
||
obf_content = self._replace_in_text(txt, fname)
|
||
|
||
results.append((clean_name, obf_content))
|
||
|
||
# Собираем ZIP + CSV
|
||
csv_str = self._build_mapping_csv()
|
||
return self._build_zip(results, csv_str), csv_str
|
||
|
||
finally:
|
||
shutil.rmtree(self._workdir, ignore_errors=True)
|
||
self._mapping.clear()
|
||
self._regex_replacements.clear()
|
||
|
||
def _expand_zips(self, files: List[Tuple[str, bytes, str]]) -> List[Tuple[str, bytes, str]]:
|
||
"""Распаковать ZIP-файлы, заменив их содержимым. Остальные файлы — как есть."""
|
||
result = []
|
||
for fname, content, ctype in files:
|
||
if fname.lower().endswith('.zip'):
|
||
try:
|
||
with zipfile.ZipFile(io.BytesIO(content)) as zf:
|
||
total_size = sum(info.file_size for info in zf.infolist())
|
||
if total_size > 500 * 1024 * 1024: # 500 MB
|
||
log.warning("ZIP too large, skipping expansion: %s", fname)
|
||
result.append((fname, content, ctype))
|
||
continue
|
||
if len(zf.infolist()) > 500:
|
||
log.warning("ZIP too many files, skipping: %s", fname)
|
||
result.append((fname, content, ctype))
|
||
continue
|
||
for info in zf.infolist():
|
||
if info.is_dir():
|
||
continue
|
||
# cp437 → utf8 (как в services/unzip.py)
|
||
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)
|
||
result.append((inner_name, inner_data, ""))
|
||
except Exception as e:
|
||
log.warning("Failed to expand ZIP %s: %s", fname, e)
|
||
result.append((fname, content, ctype))
|
||
else:
|
||
result.append((fname, content, ctype))
|
||
return result
|
||
|
||
# --- Проход 1: обнаружение ---
|
||
|
||
def _extract_text(self, fname: str, content: bytes, ctype: str) -> Tuple[str, Optional[object]]:
|
||
"""Извлечь текст из файла. Возвращает (text, docx_document_or_None)."""
|
||
doc = None
|
||
text = ""
|
||
|
||
ext = os.path.splitext(fname)[1].lower()
|
||
|
||
if ext == '.docx':
|
||
from docx import Document
|
||
doc = Document(io.BytesIO(content))
|
||
text = "\n".join(p.text for p in doc.paragraphs)
|
||
for table in doc.tables:
|
||
for row in table.rows:
|
||
text += "\n" + " | ".join(cell.text for cell in row.cells)
|
||
|
||
elif ext == '.pdf':
|
||
import pdfplumber
|
||
with pdfplumber.open(io.BytesIO(content)) as pdf:
|
||
for page in pdf.pages:
|
||
t = page.extract_text()
|
||
if t:
|
||
text += t + "\n"
|
||
for table in page.extract_tables():
|
||
for row in table:
|
||
text += "\n" + " | ".join(str(c) if c else "" for c in row)
|
||
|
||
elif ext == '.doc':
|
||
# .doc конвертируется в .docx ДО вызова — здесь уже docx
|
||
pass
|
||
|
||
else:
|
||
text = content.decode('utf-8', errors='replace')
|
||
|
||
return text, doc
|
||
|
||
def _scan_regex(self, text: str):
|
||
"""Сканировать текст regex-паттернами, заполнить словарь замен."""
|
||
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 original and original not in self._mapping:
|
||
generator = ENTITY_GENERATORS.get(entity_type, lambda x: "XXX")
|
||
self._mapping[original] = generator(original)
|
||
|
||
# Компании
|
||
for match in COMPANY_PATTERN.finditer(text):
|
||
original = match.group(0).strip()
|
||
if original and original not in self._mapping:
|
||
self._mapping[original] = generate_company(original)
|
||
|
||
def _scan_llm_ner(self, all_texts: Dict[str, str]):
|
||
"""LLM NER для обнаружения имён и адресов во всех файлах."""
|
||
combined = "\n\n---FILE---\n\n".join(
|
||
f"FILE: {fname}\n{t[:3000]}" for fname, t in all_texts.items()
|
||
)
|
||
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:
|
||
raw = self._llm_client.complete(prompt)
|
||
import json
|
||
# Игнорируем markdown-обёртку
|
||
raw = raw.strip()
|
||
if raw.startswith("```"):
|
||
raw = raw.split("\n", 1)[1]
|
||
if raw.endswith("```"):
|
||
raw = raw[:-3]
|
||
entities = json.loads(raw)
|
||
for ent in entities:
|
||
val = ent.get("value", "").strip()
|
||
if val and val not in self._mapping:
|
||
if ent.get("type") == "person":
|
||
self._mapping[val] = generate_person(val)
|
||
elif ent.get("type") == "company":
|
||
self._mapping[val] = generate_company(val)
|
||
elif ent.get("type") == "address":
|
||
self._mapping[val] = generate_address(val)
|
||
elif ent.get("type") == "passport":
|
||
self._mapping[val] = generate_passport(val)
|
||
except Exception as e:
|
||
log.warning("LLM NER failed: %s", e)
|
||
|
||
# --- Проход 2: замена ---
|
||
|
||
def _replace_in_docx(self, doc) -> bytes:
|
||
"""Заменить сущности в docx-документе (in-place)."""
|
||
# Замена в параграфах
|
||
for para in doc.paragraphs:
|
||
for run in para.runs:
|
||
run.text = self._apply_replacements(run.text)
|
||
|
||
# Замена в таблицах
|
||
for table in doc.tables:
|
||
for row in table.rows:
|
||
for cell in row.cells:
|
||
for para in cell.paragraphs:
|
||
for run in para.runs:
|
||
run.text = self._apply_replacements(run.text)
|
||
|
||
buf = io.BytesIO()
|
||
doc.save(buf)
|
||
return buf.getvalue()
|
||
|
||
def _replace_in_text(self, text: str, fname: str) -> bytes:
|
||
"""Заменить сущности в plain text (для PDF и прочих)."""
|
||
replaced = self._apply_replacements(text)
|
||
# Для PDF пока отдаём текст (MVP — без сохранения форматирования PDF)
|
||
return replaced.encode('utf-8')
|
||
|
||
def _apply_replacements(self, text: str) -> str:
|
||
"""Применить все замены из словаря mapping к строке. Сначала длинные, потом короткие."""
|
||
# Сортируем по длине оригинала (убывание) чтобы избежать частичных замен
|
||
sorted_keys = sorted(self._mapping.keys(), key=len, reverse=True)
|
||
result = text
|
||
for original in sorted_keys:
|
||
replacement = self._mapping[original]
|
||
# Только точное совпадение (не подстрока)
|
||
result = re.sub(re.escape(original), replacement, result)
|
||
return result
|
||
|
||
# --- Сборка выдачи ---
|
||
|
||
def _build_zip(self, files: List[Tuple[str, bytes]], mapping_csv: str = "") -> bytes:
|
||
"""Собрать ZIP с обфусцированными файлами + mapping.csv."""
|
||
buf = io.BytesIO()
|
||
with zipfile.ZipFile(buf, 'w', zipfile.ZIP_DEFLATED) as zf:
|
||
for fname, content in files:
|
||
zf.writestr(fname, content)
|
||
if mapping_csv:
|
||
zf.writestr("mapping.csv", mapping_csv.encode('utf-8'))
|
||
return buf.getvalue()
|
||
|
||
def _build_mapping_csv(self) -> str:
|
||
"""Собрать mapping.csv."""
|
||
buf = io.StringIO()
|
||
writer = csv.writer(buf)
|
||
writer.writerow(["тип_данных", "оригинал", "замена"])
|
||
for original, replacement in sorted(self._mapping.items()):
|
||
# Определяем тип по паттерну
|
||
etype = "text"
|
||
for t, pat in ENTITY_PATTERNS.items():
|
||
if re.match(pat, original, re.IGNORECASE):
|
||
etype = t
|
||
break
|
||
if COMPANY_PATTERN.match(original):
|
||
etype = "company"
|
||
writer.writerow([etype, original, replacement])
|
||
return buf.getvalue()
|
||
|
||
|
||
# ═══════════════════════════════════════════
|
||
# Маппинг entity_type → генератор
|
||
# ═══════════════════════════════════════════
|
||
|
||
ENTITY_GENERATORS: Dict[str, Callable] = {
|
||
"phone": generate_phone,
|
||
"email": generate_email,
|
||
"inn_ul": generate_inn10,
|
||
"inn_fl": generate_inn12,
|
||
"ogrn": generate_ogrn,
|
||
"kpp": generate_kpp,
|
||
"bik": generate_bik,
|
||
"rs": generate_rs,
|
||
"ks": generate_ks,
|
||
"passport": generate_passport,
|
||
}
|
||
|
||
|
||
def obfuscate_files(files: List[Tuple[str, bytes, str]], llm_client=None) -> Tuple[bytes, str]:
|
||
"""Удобная функция: обфусцировать список файлов → (zip_bytes, csv_string)."""
|
||
obf = TwoPassObfuscator(llm_client=llm_client)
|
||
return obf.obfuscate(files)
|