RepinoidandCopilot
1b0bcd298a
Expose extraction, analysis, and replacement progress in the file table; preserve row mapping for duplicates and collapsed archives. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
327 lines
17 KiB
Python
327 lines
17 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
|
||
from .names import resolve_name
|
||
|
||
log = logging.getLogger("drhider")
|
||
|
||
|
||
class LLMIncompleteError(Exception):
|
||
"""Processing cannot safely return files when one or more LLM chunks failed."""
|
||
|
||
def __init__(self, llm_errors: int):
|
||
self.llm_errors = llm_errors
|
||
super().__init__("LLM processing incomplete")
|
||
|
||
|
||
def _dedupe_file_names(
|
||
files: List[Tuple[str, bytes, str]]
|
||
) -> Tuple[List[Tuple[str, bytes, str]], List[List[int]]]:
|
||
"""Уникализировать имена файлов перед обработкой.
|
||
|
||
Точный дубль (одинаковое имя + одинаковое содержимое) пропускается.
|
||
При коллизии имени с ДРУГИМ содержимым добавляется суффикс _2, _3, ...
|
||
|
||
Args:
|
||
files: [(filename, content_bytes, content_type), ...]
|
||
|
||
Returns:
|
||
Новый список с уникальными именами (порядок сохранён) и исходные индексы,
|
||
соответствующие каждому файлу. Точные дубликаты ссылаются на один результат.
|
||
"""
|
||
seen: Dict[str, bytes] = {}
|
||
out: List[Tuple[str, bytes, str]] = []
|
||
source_indices: List[List[int]] = []
|
||
source_to_output: Dict[Tuple[str, bytes], int] = {}
|
||
|
||
for source_idx, (fname, content, ctype) in enumerate(files):
|
||
candidate = resolve_name(fname, content, seen)
|
||
if candidate is not None:
|
||
out.append((candidate, content, ctype))
|
||
source_to_output[(fname, content)] = len(out) - 1
|
||
source_indices.append([source_idx])
|
||
else:
|
||
source_indices[source_to_output[(fname, content)]].append(source_idx)
|
||
|
||
return out, source_indices
|
||
|
||
|
||
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._compiled_re = None # кэш скомпилированного regex (A2)
|
||
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, float, List[int]], None]] = None,
|
||
cancel_event=None, file_progress=None
|
||
) -> Tuple[bytes, str, Optional[dict]]:
|
||
"""Обфусцировать список файлов.
|
||
|
||
Все форматы → Markdown → замена → результат.
|
||
|
||
Args:
|
||
files: [(filename, content_bytes, content_type), ...]
|
||
progress_cb: Опциональный коллбек (phase, idx, total, fname, elapsed, source_idxs),
|
||
phase — extract_start/extract_done/extract_failed или
|
||
replace_start/replace_done/replace_failed; idx — 0-based индекс,
|
||
elapsed — время обработки конкретного файла (сек),
|
||
source_idxs — индексы строк исходного списка, включая дубликаты.
|
||
cancel_event: threading.Event — мягкая остановка. Проверяется между
|
||
файлами LLM и между файлами замены; текущий файл добирается.
|
||
file_progress: callable(event, fname, **fields) — прогресс по файлам/чанкам.
|
||
События: extract_complete{total_chars},
|
||
file_start{chars,chunks,idx,source_idxs},
|
||
file_chunk{chunks_done,chunks_total,idx,source_idxs},
|
||
file_done{elapsed,idx,source_idxs}.
|
||
|
||
Returns:
|
||
(zip_bytes, csv_string, meta):
|
||
zip_bytes — ZIP-архив с обфусцированными .md файлами (без mapping.csv)
|
||
csv_string — содержимое mapping.csv как строка
|
||
meta — None при штатном завершении; {"cancelled": True, "processed", "total"}
|
||
при прерывании (частичный результат)
|
||
"""
|
||
# ── Предобработка: распаковать ZIP + уникализировать имена ──
|
||
files = extractor.expand_zips(files)
|
||
files, source_indices = _dedupe_file_names(files)
|
||
|
||
# имя файла (внутр. ключ all_texts) -> idx в files (для событий file_progress)
|
||
name_to_idx = {fname: i for i, (fname, _, _) in enumerate(files)}
|
||
llm_done: set = set() # файлы, чей LLM-анализ ПОЛНОСТЬЮ завершён
|
||
|
||
def _fp(event, fname, **fields):
|
||
if event == "file_done":
|
||
llm_done.add(fname)
|
||
if file_progress:
|
||
idx = name_to_idx.get(fname)
|
||
fields["idx"] = idx
|
||
fields["source_idxs"] = source_indices[idx] if idx is not None else []
|
||
file_progress(event, fname, **fields)
|
||
|
||
try:
|
||
# ── Проход 1: сбор сущностей ──
|
||
# Извлекаем Markdown из каждого файла
|
||
all_texts: Dict[str, str] = {}
|
||
total = len(files)
|
||
file_times: List[float] = [0.0] * total
|
||
skipped: set = set() # имена файлов, которые не удалось обработать (битые)
|
||
|
||
# A1: .doc конвертируются через HTTP-сервис liberta ПОСЛЕДОВАТЕЛЬНО.
|
||
# Без параллельности: liberta не выдерживает конкурентные вызовы
|
||
# (libreoffice падает с javaldx).
|
||
|
||
cancelled_extract = False
|
||
for i, (fname, content, ctype) in enumerate(files):
|
||
# Мягкая остановка в фазе извлечения (между файлами) — не ждём все .doc
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
log.info("obfuscate: cancelled during extract, processed=%d/%d", i, total)
|
||
cancelled_extract = True
|
||
break
|
||
display_name = os.path.basename(fname) or fname
|
||
if progress_cb:
|
||
progress_cb("extract_start", i, total, display_name, 0.0, source_indices[i])
|
||
|
||
t0 = time.time()
|
||
try:
|
||
text = extractor.extract_text(fname, content, ctype)
|
||
except Exception as e:
|
||
log.warning("obfuscate: skip file %r: %r", fname, e)
|
||
skipped.add(fname)
|
||
if progress_cb:
|
||
progress_cb("extract_failed", i, total, display_name, 0.0, source_indices[i])
|
||
continue
|
||
all_texts[fname] = text
|
||
if progress_cb:
|
||
progress_cb("extract_done", i, total, display_name, 0.0, source_indices[i])
|
||
|
||
# Regex-сканирование (быстрое, локальное)
|
||
if text and not text.startswith("[DOC binary"):
|
||
scanner.scan_regex(text, self._mapping, self._counters)
|
||
file_times[i] += time.time() - t0
|
||
|
||
# Объём извлечённого текста — для оценки времени (глобальной ETA)
|
||
total_chars = sum(len(t) for t in all_texts.values())
|
||
if file_progress:
|
||
file_progress("extract_complete", None, total_chars=total_chars,
|
||
per_file={fname: len(t) for fname, t in all_texts.items()})
|
||
|
||
# LLM-сканирование (получает уже найденное regex'ом чтобы не дублировать)
|
||
cancelled_llm = cancelled_extract
|
||
if not cancelled_llm and self._llm_client:
|
||
try:
|
||
llm_failures = scanner.scan_llm_ner(
|
||
all_texts, self._mapping, self._llm_client, self._counters,
|
||
cancel_event=cancel_event, file_progress=_fp,
|
||
)
|
||
except scanner.CancelRequested as e:
|
||
if e.llm_failures:
|
||
raise LLMIncompleteError(e.llm_failures)
|
||
log.info("obfuscate: LLM cancelled, llm_done=%d", len(llm_done))
|
||
cancelled_llm = True
|
||
else:
|
||
if llm_failures:
|
||
raise LLMIncompleteError(llm_failures)
|
||
|
||
# Предсортировать ключи один раз (по убыванию длины)
|
||
self._sorted_keys = sorted(
|
||
self._mapping.keys(), key=len, reverse=True
|
||
)
|
||
# A2: кэш скомпилированного regex — один на все файлы прохода 2
|
||
self._compiled_re = replacer._build_combined_re(self._sorted_keys)
|
||
|
||
# ── Проход 2: замена сущностей ──
|
||
cancelled = cancelled_llm
|
||
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
|
||
|
||
# Битый файл (не удалось извлечь текст) — пропускаем
|
||
if fname in skipped:
|
||
continue
|
||
|
||
# При прерывании в LLM — в результат идут ТОЛЬКО файлы с завершённым LLM
|
||
if cancelled_llm and fname not in llm_done:
|
||
continue
|
||
|
||
# Мягкая остановка в фазе замены (между файлами).
|
||
# НЕ проверяем, если LLM уже прерван: файлы из llm_done добираем (сохраняем всё готовое).
|
||
if not cancelled_llm and cancel_event is not None and cancel_event.is_set():
|
||
log.info("obfuscate: cancelled in replace phase, processed=%d", len(results))
|
||
cancelled = True
|
||
break
|
||
|
||
if progress_cb:
|
||
progress_cb("replace_start", i, total, display_name, 0.0, source_indices[i])
|
||
t0 = time.time()
|
||
obf_content = content # По умолчанию — без изменений
|
||
|
||
if fname.lower().endswith('.csv'):
|
||
try:
|
||
obf_content = extractor.obfuscate_csv(
|
||
content, self._mapping, self._sorted_keys, self._compiled_re
|
||
)
|
||
except Exception as e:
|
||
log.warning("obfuscate: skip CSV %r during replacement: %r", fname, e)
|
||
skipped.add(fname)
|
||
if progress_cb:
|
||
progress_cb("replace_failed", i, total, display_name, 0.0, source_indices[i])
|
||
continue
|
||
elif fname.endswith('.doc'):
|
||
# .doc — конвертируется через liberta, обрабатывается как .docx
|
||
md_text = all_texts.get(fname, '')
|
||
replaced = replacer.apply_replacements(
|
||
md_text, self._mapping, self._sorted_keys, self._compiled_re
|
||
)
|
||
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, self._compiled_re
|
||
)
|
||
# Имя файла: заменить расширение на .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(
|
||
"replace_done", i, total, display_name,
|
||
round(file_times[i], 2), source_indices[i],
|
||
)
|
||
|
||
# ── Сборка результата ──
|
||
csv_str = builder.build_mapping_csv(self._mapping)
|
||
zip_data = builder.build_zip(results, csv_str)
|
||
|
||
meta = None
|
||
if cancelled:
|
||
expected_total = total - len(skipped)
|
||
meta = {"cancelled": True, "processed": len(results), "total": expected_total}
|
||
return zip_data, csv_str, meta
|
||
|
||
finally:
|
||
# Очистка состояния (обфускатор может использоваться повторно)
|
||
self._mapping.clear()
|
||
self._sorted_keys.clear()
|
||
self._compiled_re = None
|
||
self._counters.clear()
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════
|
||
# Удобная функция для быстрого вызова
|
||
# ═══════════════════════════════════════════════════════════════════════
|
||
|
||
def obfuscate_files(
|
||
files: List[Tuple[str, bytes, str]], llm_client=None,
|
||
progress_cb: Optional[Callable[[str, int, int, str, float, List[int]], None]] = None,
|
||
cancel_event=None, file_progress=None
|
||
) -> Tuple[bytes, str, Optional[dict]]:
|
||
"""Обфусцировать список файлов — удобная функция.
|
||
|
||
Создаёт экземпляр TwoPassObfuscator и вызывает .obfuscate().
|
||
|
||
Args:
|
||
files: [(filename, content_bytes, content_type), ...]
|
||
llm_client: Опциональный LLM-клиент
|
||
progress_cb: Опциональный коллбек (phase, idx, total, fname, elapsed, source_idxs)
|
||
cancel_event: threading.Event — мягкая остановка (см. TwoPassObfuscator.obfuscate)
|
||
file_progress: callable(event, fname, **fields) — прогресс по файлам/чанкам
|
||
|
||
Returns:
|
||
(zip_bytes, csv_string, meta)
|
||
"""
|
||
obf = TwoPassObfuscator(llm_client=llm_client)
|
||
return obf.obfuscate(files, progress_cb=progress_cb,
|
||
cancel_event=cancel_event, file_progress=file_progress)
|