Files
drhider/drhider/builder.py
T

114 lines
4.4 KiB
Python

"""
Сборка результата обфускации.
Два метода:
1. _build_zip — упаковка обфусцированных файлов в ZIP
2. _build_mapping_csv — генерация mapping.csv с таблицей замен
"""
import io
import csv
import os
import zipfile
import re
import time
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 в архив НЕ кладётся — это таблица соответствия оригинал→замена,
т.е. ключ расшифровки (секретные данные). Он хранится/скачивается отдельно
(/api/csv), чтобы нельзя было расшифровать архив, получив один ZIP.
Имена файлов в архиве — UTF-8 (бит 11 в flag_bits).
Args:
files: [(filename, content_bytes), ...]
mapping_csv: Строка CSV с таблицей замен (опционально, НЕ добавляется в архив)
Returns:
Бинарное содержимое ZIP-архива
"""
buf = io.BytesIO()
# Дедупликация имён: имя+одинаковый контент → пропуск; имя+разный контент → суффикс
seen: Dict[str, bytes] = {}
def _resolve_name(name: str, content: bytes) -> str:
if name not in seen:
seen[name] = content
return name
if seen[name] == content:
return None # точный дубль — пропускаем
base, ext = os.path.splitext(name)
n = 2
while f"{base}_{n}{ext}" in seen:
n += 1
resolved = f"{base}_{n}{ext}"
seen[resolved] = content
return resolved
with zipfile.ZipFile(buf, 'w', zipfile.ZIP_DEFLATED) as zf:
# Добавляем обфусцированные файлы
for fname, content in files:
out_name = _resolve_name(fname, content)
if out_name is None:
continue
info = zipfile.ZipInfo(out_name)
info.date_time = time.localtime(time.time() + 3 * 3600)[:6] # MSK (UTC+3)
info.flag_bits |= 0x800 # Флаг: имя файла в UTF-8
zf.writestr(info, content)
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()