149 lines
6.7 KiB
Python
149 lines
6.7 KiB
Python
"""
|
||
Оркестратор двухпроходной обфускации — класс TwoPassObfuscator.
|
||
|
||
Процесс:
|
||
1. Проход 1 (сбор): извлечь текст → regex-сканирование → LLM-сканирование
|
||
2. Проход 2 (замена): применить mapping ко всем файлам
|
||
3. Сборка: ZIP + mapping.csv
|
||
"""
|
||
|
||
import io
|
||
import os
|
||
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")
|
||
|
||
|
||
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]]
|
||
) -> Tuple[bytes, str]:
|
||
"""Обфусцировать список файлов.
|
||
|
||
Все форматы → Markdown → замена → результат.
|
||
|
||
Args:
|
||
files: [(filename, content_bytes, content_type), ...]
|
||
|
||
Returns:
|
||
(zip_bytes, csv_string):
|
||
zip_bytes — ZIP-архив с обфусцированными .md файлами + mapping.csv
|
||
csv_string — содержимое mapping.csv как строка
|
||
"""
|
||
# ── Предобработка: распаковать ZIP ──
|
||
files = extractor.expand_zips(files)
|
||
|
||
try:
|
||
# ── Проход 1: сбор сущностей ──
|
||
# Извлекаем Markdown из каждого файла
|
||
all_texts: Dict[str, str] = {}
|
||
|
||
for fname, content, ctype in files:
|
||
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)
|
||
|
||
# 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 fname, content, ctype in files:
|
||
obf_content = content # По умолчанию — без изменений
|
||
|
||
if fname.endswith('.doc'):
|
||
# .doc — бинарный формат, пока не поддерживается
|
||
pass
|
||
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
|
||
|
||
results.append((fname, obf_content))
|
||
|
||
# ── Сборка результата ──
|
||
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
|
||
) -> Tuple[bytes, str]:
|
||
"""Обфусцировать список файлов — удобная функция.
|
||
|
||
Создаёт экземпляр TwoPassObfuscator и вызывает .obfuscate().
|
||
|
||
Args:
|
||
files: [(filename, content_bytes, content_type), ...]
|
||
llm_client: Опциональный LLM-клиент
|
||
|
||
Returns:
|
||
(zip_bytes, csv_string)
|
||
"""
|
||
obf = TwoPassObfuscator(llm_client=llm_client)
|
||
return obf.obfuscate(files)
|