211 lines
9.5 KiB
Python
211 lines
9.5 KiB
Python
"""
|
||
Оркестратор двухпроходной обфускации — класс TwoPassObfuscator.
|
||
|
||
Процесс:
|
||
1. Проход 1 (сбор): извлечь текст → regex-сканирование → LLM-сканирование
|
||
2. Проход 2 (замена): применить mapping ко всем файлам
|
||
3. Сборка: ZIP + mapping.csv
|
||
"""
|
||
|
||
import io
|
||
import os
|
||
import time
|
||
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")
|
||
|
||
|
||
def _dedupe_file_names(
|
||
files: List[Tuple[str, bytes, str]]
|
||
) -> List[Tuple[str, bytes, str]]:
|
||
"""Уникализировать имена файлов перед обработкой.
|
||
|
||
Точный дубль (одинаковое имя + одинаковое содержимое) пропускается.
|
||
При коллизии имени с ДРУГИМ содержимым добавляется суффикс _2, _3, ...
|
||
|
||
Args:
|
||
files: [(filename, content_bytes, content_type), ...]
|
||
|
||
Returns:
|
||
Новый список с уникальными именами (порядок сохранён).
|
||
"""
|
||
seen: Dict[str, bytes] = {}
|
||
out: List[Tuple[str, bytes, str]] = []
|
||
|
||
for fname, content, ctype in files:
|
||
candidate = fname
|
||
if candidate in seen:
|
||
if seen[candidate] == content:
|
||
continue # точный дубль — пропускаем
|
||
base, ext = os.path.splitext(fname)
|
||
n = 2
|
||
while f"{base}_{n}{ext}" in seen:
|
||
n += 1
|
||
candidate = f"{base}_{n}{ext}"
|
||
seen[candidate] = content
|
||
out.append((candidate, content, ctype))
|
||
|
||
return out
|
||
|
||
|
||
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._counters: Dict[str, int] = {} # Глобальные счётчики токенов
|
||
self._llm_client = llm_client
|
||
|
||
# ═══════════════════════════════════════════════════════════════════
|
||
# Главная точка входа
|
||
# ═══════════════════════════════════════════════════════════════════
|
||
|
||
def obfuscate(
|
||
self, files: List[Tuple[str, bytes, str]],
|
||
progress_cb: Optional[Callable[[str, int, int, str], None]] = None
|
||
) -> Tuple[bytes, str]:
|
||
"""Обфусцировать список файлов.
|
||
|
||
Все форматы → Markdown → замена → результат.
|
||
|
||
Args:
|
||
files: [(filename, content_bytes, content_type), ...]
|
||
progress_cb: Опциональный коллбек (phase, idx, total, fname, elapsed),
|
||
где phase ∈ {"start", "done"}, idx — 0-based индекс,
|
||
elapsed — время обработки конкретного файла (сек).
|
||
|
||
Returns:
|
||
(zip_bytes, csv_string):
|
||
zip_bytes — ZIP-архив с обфусцированными .md файлами + mapping.csv
|
||
csv_string — содержимое mapping.csv как строка
|
||
"""
|
||
# ── Предобработка: распаковать ZIP + уникализировать имена ──
|
||
files = extractor.expand_zips(files)
|
||
files = _dedupe_file_names(files)
|
||
|
||
try:
|
||
# ── Проход 1: сбор сущностей ──
|
||
# Извлекаем Markdown из каждого файла
|
||
all_texts: Dict[str, str] = {}
|
||
total = len(files)
|
||
file_times: List[float] = [0.0] * total
|
||
|
||
for i, (fname, content, ctype) in enumerate(files):
|
||
display_name = os.path.basename(fname) or fname
|
||
if progress_cb:
|
||
progress_cb("start", i, total, display_name, 0.0)
|
||
|
||
t0 = time.time()
|
||
text = extractor.extract_text(fname, content, ctype)
|
||
all_texts[fname] = text
|
||
|
||
# Regex-сканирование (быстрое, локальное)
|
||
if text and not text.startswith("[DOC binary"):
|
||
scanner.scan_regex(text, self._mapping, self._counters)
|
||
file_times[i] += time.time() - t0
|
||
|
||
# LLM-сканирование (получает уже найденное regex'ом чтобы не дублировать)
|
||
if self._llm_client:
|
||
scanner.scan_llm_ner(all_texts, self._mapping, self._llm_client, self._counters)
|
||
|
||
# Предсортировать ключи один раз (по убыванию длины)
|
||
self._sorted_keys = sorted(
|
||
self._mapping.keys(), key=len, reverse=True
|
||
)
|
||
|
||
# ── Проход 2: замена сущностей ──
|
||
results: List[Tuple[str, bytes]] = []
|
||
|
||
for i, (in_fname, content, ctype) in enumerate(files):
|
||
fname = in_fname
|
||
display_name = os.path.basename(in_fname) or in_fname
|
||
|
||
t0 = time.time()
|
||
obf_content = content # По умолчанию — без изменений
|
||
|
||
if fname.endswith('.doc'):
|
||
# .doc — конвертируется через liberta, обрабатывается как .docx
|
||
md_text = all_texts.get(fname, '')
|
||
replaced = replacer.apply_replacements(
|
||
md_text, self._mapping, self._sorted_keys
|
||
)
|
||
md_name = os.path.splitext(fname)[0] + '.md'
|
||
obf_content = replaced.encode('utf-8')
|
||
fname = md_name
|
||
else:
|
||
# Все форматы → Markdown → замена в тексте
|
||
md_text = all_texts.get(fname, '')
|
||
replaced = replacer.apply_replacements(
|
||
md_text, self._mapping, self._sorted_keys
|
||
)
|
||
# Имя файла: заменить расширение на .md
|
||
md_name = os.path.splitext(fname)[0] + '.md'
|
||
obf_content = replaced.encode('utf-8')
|
||
fname = md_name
|
||
|
||
file_times[i] += time.time() - t0
|
||
results.append((fname, obf_content))
|
||
if progress_cb:
|
||
progress_cb("done", i, total, display_name, round(file_times[i], 2))
|
||
|
||
# ── Сборка результата ──
|
||
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()
|
||
self._counters.clear()
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════
|
||
# Удобная функция для быстрого вызова
|
||
# ═══════════════════════════════════════════════════════════════════════
|
||
|
||
def obfuscate_files(
|
||
files: List[Tuple[str, bytes, str]], llm_client=None,
|
||
progress_cb: Optional[Callable[[str, int, int, str], None]] = None
|
||
) -> Tuple[bytes, str]:
|
||
"""Обфусцировать список файлов — удобная функция.
|
||
|
||
Создаёт экземпляр TwoPassObfuscator и вызывает .obfuscate().
|
||
|
||
Args:
|
||
files: [(filename, content_bytes, content_type), ...]
|
||
llm_client: Опциональный LLM-клиент
|
||
progress_cb: Опциональный коллбек (phase, idx, total, fname)
|
||
|
||
Returns:
|
||
(zip_bytes, csv_string)
|
||
"""
|
||
obf = TwoPassObfuscator(llm_client=llm_client)
|
||
return obf.obfuscate(files, progress_cb=progress_cb)
|