From 94e588aaba91f57b7d5a1bf470e1ae6f03793329 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CNaeel=E2=80=9D?= Date: Tue, 18 Aug 2026 23:03:19 +0400 Subject: [PATCH] =?UTF-8?q?v0.0.32:=20=D1=84=D0=B8=D0=BA=D1=81=D1=8B=20?= =?UTF-8?q?=D0=BF=D0=BE=20=D1=80=D0=B5=D0=B2=D1=8C=D1=8E=20=E2=80=94=20?= =?UTF-8?q?=D0=B5=D0=B4=D0=B8=D0=BD=D1=8B=D0=B9=20=D0=BE=D0=B1=D1=84=D1=83?= =?UTF-8?q?=D1=81=D0=BA=D0=B0=D1=82=D0=BE=D1=80,=20=D1=80=D0=B5=D0=BA?= =?UTF-8?q?=D1=83=D1=80=D1=81=D0=B8=D1=8F=20ZIP,=20=D0=BA=D0=B8=D1=80?= =?UTF-8?q?=D0=B8=D0=BB=D0=BB=D0=B8=D1=86=D0=B0=201=D0=A1,=20=D0=B4=D0=B5?= =?UTF-8?q?=D0=B4=D1=83=D0=BF,=20SSE-=D0=B4=D0=B8=D1=81=D0=BA=D0=BE=D0=BD?= =?UTF-8?q?=D0=BD=D0=B5=D0=BA=D1=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- drhider/builder.py | 23 ++++++- drhider/extractor.py | 78 ++++++++++++--------- drhider/obfuscator.py | 58 ++++++++++++++-- site/app.py | 2 +- site/routes/api_bp.py | 156 +++++++++++++++++++++++++----------------- site/session.py | 6 ++ 6 files changed, 220 insertions(+), 103 deletions(-) diff --git a/drhider/builder.py b/drhider/builder.py index 4eb8611..c2b5ac2 100644 --- a/drhider/builder.py +++ b/drhider/builder.py @@ -8,6 +8,7 @@ import io import csv +import os import zipfile import re import time @@ -34,10 +35,30 @@ def build_zip(files: List[Tuple[str, bytes]], mapping_csv: str = "") -> bytes: """ 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: - info = zipfile.ZipInfo(fname) + 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) diff --git a/drhider/extractor.py b/drhider/extractor.py index ed23f9c..ee8d396 100644 --- a/drhider/extractor.py +++ b/drhider/extractor.py @@ -247,61 +247,73 @@ def expand_zips(files: List[Tuple[str, bytes, str]]) -> List[Tuple[str, bytes, s Returns: Новый список файлов (ZIP раскрыты, остальные как есть) """ - result: List[Tuple[str, bytes, str]] = [] + # Лимиты защиты от ZIP-бомб (глобально на весь вызов) + MAX_FILES_IN_ARCHIVE = 500 + MAX_UNCOMPRESSED = 500 * 1024 * 1024 # 500 MB + MAX_RATIO = 100 + + result: List[Tuple[str, bytes, str]] = [] + queue: List[Tuple[str, bytes, str]] = list(files) + total_uncompressed = 0 + + def _decode_name(name: str, info) -> str: + """Декодировать имя из ZIP. + + Если флаг UTF-8 (bit 11) выставлен — имя уже корректное (берём как есть). + Иначе zipfile декодировал имя как CP437. Для кириллицы из 1С (CP866) + перекодируем CP437→CP866; при неудаче оставляем как есть. + """ + if not (info.flag_bits & 0x800) and any(ord(c) > 127 for c in name): + try: + return name.encode("cp437").decode("cp866") + except (UnicodeDecodeError, UnicodeEncodeError): + pass + return name + + while queue: + fname, content, ctype = queue.pop(0) - for fname, content, ctype in files: if not fname.lower().endswith('.zip'): result.append((fname, content, ctype)) continue try: with zipfile.ZipFile(io.BytesIO(content)) as zf: - if len(zf.infolist()) > 500: + infos = [i for i in zf.infolist() if not i.is_dir()] + + if len(infos) > MAX_FILES_IN_ARCHIVE: log.warning("ZIP too many files, skipping: %s", fname) result.append((fname, content, ctype)) continue - total_uncompressed = 0 - - for info in zf.infolist(): - if info.is_dir(): - continue - + # Проверяем лимиты ДО чтения содержимого (защита от бомб) + bomb = False + for info in infos: if info.compress_size > 0: ratio = info.file_size / info.compress_size - if ratio > 100: - log.warning( - "ZIP bomb ratio %.0f:1, skipping: %s", ratio, fname - ) - result.append((fname, content, ctype)) + if ratio > MAX_RATIO: + bomb = True break + if total_uncompressed + info.file_size > MAX_UNCOMPRESSED: + bomb = True + break - name = info.filename - try: - name = name.encode("cp437").decode("utf-8", errors="replace") - except (UnicodeDecodeError, UnicodeEncodeError): - pass + if bomb: + log.warning("ZIP bomb/limit, skipping: %s", fname) + result.append((fname, content, ctype)) + continue + for info in infos: + name = _decode_name(info.filename, info) name = os.path.basename(name) - if ( - not name - or name.endswith("/") - or ".." in name - or "/" in name - or "\\" in name - ): + if not name: continue inner_data = zf.read(info) total_uncompressed += len(inner_data) - if total_uncompressed > 500 * 1024 * 1024: - log.warning( - "ZIP uncompressed limit exceeded, stopping: %s", fname - ) - break - - result.append((name, inner_data, "")) + # Вложенные ZIP добавляем в очередь на повторную распаковку + queue.append((name, inner_data, "")) except Exception as e: log.warning("Failed to expand ZIP %s: %s", fname, e) diff --git a/drhider/obfuscator.py b/drhider/obfuscator.py index e7da7d3..de3936e 100644 --- a/drhider/obfuscator.py +++ b/drhider/obfuscator.py @@ -20,6 +20,39 @@ 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: """Двухпроходный обфускатор документов. @@ -52,7 +85,8 @@ class TwoPassObfuscator: # ═══════════════════════════════════════════════════════════════════ def obfuscate( - self, files: List[Tuple[str, bytes, str]] + self, files: List[Tuple[str, bytes, str]], + progress_cb: Optional[Callable[[str, int, int, str], None]] = None ) -> Tuple[bytes, str]: """Обфусцировать список файлов. @@ -60,14 +94,18 @@ class TwoPassObfuscator: Args: files: [(filename, content_bytes, content_type), ...] + progress_cb: Опциональный коллбек (phase, idx, total, fname), + где phase ∈ {"start", "done"}, idx — 0-based индекс. + Вызывается вокруг финальной обработки каждого файла. Returns: (zip_bytes, csv_string): zip_bytes — ZIP-архив с обфусцированными .md файлами + mapping.csv csv_string — содержимое mapping.csv как строка """ - # ── Предобработка: распаковать ZIP ── + # ── Предобработка: распаковать ZIP + уникализировать имена ── files = extractor.expand_zips(files) + files = _dedupe_file_names(files) try: # ── Проход 1: сбор сущностей ── @@ -93,8 +131,14 @@ class TwoPassObfuscator: # ── Проход 2: замена сущностей ── results: List[Tuple[str, bytes]] = [] + total = len(files) + + for i, (in_fname, content, ctype) in enumerate(files): + fname = in_fname + display_name = os.path.basename(in_fname) or in_fname + if progress_cb: + progress_cb("start", i, total, display_name) - for fname, content, ctype in files: obf_content = content # По умолчанию — без изменений if fname.endswith('.doc'): @@ -118,6 +162,8 @@ class TwoPassObfuscator: fname = md_name results.append((fname, obf_content)) + if progress_cb: + progress_cb("done", i, total, display_name) # ── Сборка результата ── csv_str = builder.build_mapping_csv(self._mapping) @@ -137,7 +183,8 @@ class TwoPassObfuscator: # ═══════════════════════════════════════════════════════════════════════ def obfuscate_files( - files: List[Tuple[str, bytes, str]], llm_client=None + files: List[Tuple[str, bytes, str]], llm_client=None, + progress_cb: Optional[Callable[[str, int, int, str], None]] = None ) -> Tuple[bytes, str]: """Обфусцировать список файлов — удобная функция. @@ -146,9 +193,10 @@ def obfuscate_files( 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) + return obf.obfuscate(files, progress_cb=progress_cb) diff --git a/site/app.py b/site/app.py index d37f67e..d3a5e44 100644 --- a/site/app.py +++ b/site/app.py @@ -20,7 +20,7 @@ if _sys_path_root not in sys.path: sys.path.insert(0, _sys_path_root) # Версия приложения (меняется при изменениях) -VERSION = "0.0.31" +VERSION = "0.0.32" def create_app(): diff --git a/site/routes/api_bp.py b/site/routes/api_bp.py index 7194d7f..bc993d4 100644 --- a/site/routes/api_bp.py +++ b/site/routes/api_bp.py @@ -11,84 +11,130 @@ Five endpoints: import io import json +import queue +import threading import zipfile import traceback from datetime import datetime, timedelta from flask import Blueprint, request, send_file, jsonify, Response, stream_with_context from drhider import obfuscate_files, LLMClient -from drhider.builder import build_zip, build_mapping_csv from session import (create_session, add_file, get_files, store_result, get_result, store_csv, get_csv, cleanup, file_count) api_bp = Blueprint("api", __name__, url_prefix="/api") +def _disconnect_exceptions(): + """Исключения, означающие отключение клиента SSE.""" + return (GeneratorExit, BrokenPipeError, ConnectionResetError) + + @api_bp.route("/upload", methods=["POST"]) def upload(): - """Upload one file to session.""" + """Upload files to session (один или несколько).""" sid = request.form.get("session", "") if not sid: sid = create_session() uploaded = request.files.getlist("files") if not uploaded: return jsonify({"ok": False, "error": "No file"}), 400 - f = uploaded[0] - if not f.filename: - return jsonify({"ok": False, "error": "No filename"}), 400 - ok = add_file(sid, f.filename, f.read()) - if not ok: - return jsonify({"ok": False, "error": "Session not found"}), 404 + + added = 0 + had_unnamed = False + for f in uploaded: + if not f.filename: + had_unnamed = True + continue + if not add_file(sid, f.filename, f.read()): + return jsonify({"ok": False, "error": "Session not found"}), 404 + added += 1 + + if added == 0: + err = "No filename" if had_unnamed else "No file" + return jsonify({"ok": False, "error": err}), 400 + return jsonify({"ok": True, "session": sid, "count": file_count(sid)}) @api_bp.route("/process_stream/", methods=["GET"]) def process_stream(sid): - """SSE: process all session files one by one, streaming per-file progress.""" + """SSE: process all session files, streaming per-file progress. + + Все файлы обрабатываются ЕДИНЫМ вызовом obfuscate_files (общий mapping, + согласованные токены). Обработка идёт в отдельном потоке; прогресс + передаётся через очередь. Разрыв соединения клиента корректно + перехватывается и останавливает генератор. + """ files = get_files(sid) if files is None: return jsonify({"ok": False, "error": "Session not found"}), 404 if not files: return jsonify({"ok": False, "error": "No files"}), 400 - total = len(files) + all_files = [(fname, content, "") for fname, content in files] def generate(): llm = LLMClient() - all_results = [] - all_mapping = {} - for idx, (fname, content) in enumerate(files): - # Отправляем: начали файл - try: - yield f"event: start\ndata: {json.dumps({'idx': idx, 'name': fname, 'total': total})}\n\n" - except GeneratorExit: - return # клиент отключился + q = queue.Queue() + cancel = threading.Event() - zip_data, _csv = obfuscate_files([(fname, content, "")], llm_client=llm) - with zipfile.ZipFile(io.BytesIO(zip_data)) as zf: - for name in zf.namelist(): - if name == "mapping.csv": - csv_text = zf.read(name).decode("utf-8") - lines = csv_text.strip().split("\n") - if len(lines) > 1: - for line in lines[1:]: - parts = line.split(",", 2) - if len(parts) >= 3: - all_mapping[parts[1]] = parts[2] - elif not name.endswith("/"): - all_results.append((name, zf.read(name))) - # Отправляем: закончили файл (если клиент ушёл — освобождаем поток) - try: - yield f"event: done\ndata: {json.dumps({'idx': idx, 'name': fname, 'total': total})}\n\n" - except GeneratorExit: - return # клиент отключился, НЕ обрабатываем остальные файлы + def progress(phase, idx, total_, name): + q.put(("progress", phase, idx, name, total_)) - csv_str = build_mapping_csv(all_mapping) if all_mapping else "" - final_zip = build_zip(all_results) - store_result(sid, final_zip) - if csv_str: - store_csv(sid, csv_str) - yield f"event: complete\ndata: {json.dumps({'total': len(all_results)})}\n\n" + def worker(): + try: + zip_data, csv_str = obfuscate_files( + all_files, llm_client=llm, progress_cb=progress + ) + q.put(("result", zip_data, csv_str)) + except Exception as e: + q.put(("error", repr(e))) + + threading.Thread(target=worker, daemon=True).start() + + while True: + try: + evt = q.get(timeout=1) + except queue.Empty: + if cancel.is_set(): + return + continue + + kind = evt[0] + + if kind == "progress": + _, phase, idx, name, total_ = evt + try: + yield ( + f"event: {phase}\n" + f"data: {json.dumps({'idx': idx, 'name': name, 'total': total_})}\n\n" + ) + except _disconnect_exceptions(): + cancel.set() + return + + elif kind == "result": + _, zip_data, csv_str = evt + store_result(sid, zip_data) + if csv_str: + store_csv(sid, csv_str) + count = 0 + with zipfile.ZipFile(io.BytesIO(zip_data)) as zf: + count = len([n for n in zf.namelist() if n != "mapping.csv"]) + try: + yield f"event: complete\ndata: {json.dumps({'total': count})}\n\n" + except _disconnect_exceptions(): + return + return + + elif kind == "error": + _, msg = evt + try: + yield f"event: error\ndata: {json.dumps({'error': msg})}\n\n" + except _disconnect_exceptions(): + return + return return Response( stream_with_context(generate()), @@ -99,7 +145,7 @@ def process_stream(sid): @api_bp.route("/process/", methods=["POST"]) def process(sid): - """Process all session files one by one -> ZIP.""" + """Process all session files -> ZIP (единым вызовом, общий mapping).""" files = get_files(sid) if files is None: return jsonify({"ok": False, "error": "Session not found"}), 404 @@ -107,28 +153,12 @@ def process(sid): return jsonify({"ok": False, "error": "No files"}), 400 try: llm = LLMClient() - all_results = [] - all_mapping = {} - for fname, content in files: - zip_data, _csv = obfuscate_files([(fname, content, "")], llm_client=llm) - with zipfile.ZipFile(io.BytesIO(zip_data)) as zf: - for name in zf.namelist(): - if name == "mapping.csv": - csv_text = zf.read(name).decode("utf-8") - lines = csv_text.strip().split("\n") - if len(lines) > 1: - for line in lines[1:]: - parts = line.split(",", 2) - if len(parts) >= 3: - all_mapping[parts[1]] = parts[2] - elif not name.endswith("/"): - all_results.append((name, zf.read(name))) - csv_str = build_mapping_csv(all_mapping) if all_mapping else "" - final_zip = build_zip(all_results) # CSV — отдельно, не в ZIP - store_result(sid, final_zip) + all_files = [(fname, content, "") for fname, content in files] + zip_data, csv_str = obfuscate_files(all_files, llm_client=llm) + store_result(sid, zip_data) if csv_str: store_csv(sid, csv_str) - return jsonify({"ok": True, "status": "done", "files": len(all_results)}) + return jsonify({"ok": True, "status": "done"}) except Exception as e: traceback.print_exc() return jsonify({"ok": False, "error": str(e)}), 500 diff --git a/site/session.py b/site/session.py index 66e1822..7f5ca8f 100644 --- a/site/session.py +++ b/site/session.py @@ -21,6 +21,9 @@ from typing import Dict, List, Tuple, Optional TTL_SECONDS = 30 * 60 # 30 минут +# Максимальный суммарный объём файлов в одной сессии (защита памяти) +MAX_SESSION_BYTES = 500 * 1024 * 1024 # 500 MB + _sessions: Dict[str, dict] = {} _lock = threading.Lock() @@ -73,6 +76,9 @@ def add_file(sid: str, filename: str, content: bytes) -> bool: s = _sessions.get(sid) if not s: return False + total = sum(len(c) for _, c in s["files"]) + if total + len(content) > MAX_SESSION_BYTES: + return False # превышен суммарный лимит сессии s["files"].append((filename, content)) return True