diff --git a/site/app.py b/site/app.py index 7c3e7b1..097864e 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.1" +VERSION = "0.0.2" def create_app(): @@ -51,6 +51,7 @@ def create_app(): response.headers["Cache-Control"] = "no-cache, no-store, must-revalidate" response.headers["Pragma"] = "no-cache" response.headers["Expires"] = "0" + response.headers["Connection"] = "close" return response return app diff --git a/site/routes/__init__.py b/site/routes/__init__.py index 5f23f5c..a475b5b 100644 --- a/site/routes/__init__.py +++ b/site/routes/__init__.py @@ -1,20 +1,13 @@ -""" -Регистрация всех blueprint'ов приложения. - -Импортируется из site/app.py при создании Flask-приложения. -""" - from .main_bp import main_bp from .health_bp import health_bp -from .api_bp import api_bp +from .upload_bp import upload_bp +from .process_bp import process_bp +from .download_bp import download_bp def register_routes(app): - """Зарегистрировать все blueprint'ы на Flask-приложении. - - Args: - app: Экземпляр Flask-приложения - """ app.register_blueprint(main_bp) app.register_blueprint(health_bp) - app.register_blueprint(api_bp) + app.register_blueprint(upload_bp) + app.register_blueprint(process_bp) + app.register_blueprint(download_bp) diff --git a/site/routes/api_bp.py b/site/routes/api_bp.py deleted file mode 100644 index a32b301..0000000 --- a/site/routes/api_bp.py +++ /dev/null @@ -1,83 +0,0 @@ -""" -Blueprint: API DrHider. - -Three endpoints: -- POST /api/upload — upload one file -> {session_id} -- POST /api/process/{sid} — process all session files -> {status:"done"} -- GET /api/download/{sid} — download ZIP -""" - -import io -import zipfile -import traceback -from flask import Blueprint, request, send_file, jsonify - -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, cleanup, file_count - -api_bp = Blueprint("api", __name__, url_prefix="/api") - - -@api_bp.route("/upload", methods=["POST"]) -def upload(): - """Upload one file 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 - return jsonify({"ok": True, "session": sid, "count": file_count(sid)}) - - -@api_bp.route("/process/", methods=["POST"]) -def process(sid): - """Process all session files one by one -> ZIP.""" - 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 - 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, mapping_csv=csv_str) - store_result(sid, final_zip) - return jsonify({"ok": True, "status": "done", "files": len(all_results)}) - except Exception as e: - traceback.print_exc() - return jsonify({"ok": False, "error": str(e)}), 500 - - -@api_bp.route("/download/", methods=["GET"]) -def download(sid): - """Download result and cleanup session.""" - zip_data = get_result(sid) - if zip_data is None: - return jsonify({"ok": False, "error": "Not found"}), 404 - cleanup(sid) - return send_file(io.BytesIO(zip_data), mimetype="application/zip", - as_attachment=True, download_name="drhider_output.zip") diff --git a/site/routes/download_bp.py b/site/routes/download_bp.py new file mode 100644 index 0000000..623b6b6 --- /dev/null +++ b/site/routes/download_bp.py @@ -0,0 +1,21 @@ +""" +GET /api/download/{sid} — скачать ZIP и удалить сессию. + +После этого вызова — всё чисто, памяти ноль. +""" + +import io +from flask import Blueprint, send_file, jsonify +from session import get_result, cleanup + +download_bp = Blueprint("download", __name__, url_prefix="/api") + + +@download_bp.route("/download/", methods=["GET"]) +def download(sid): + zip_data = get_result(sid) + if zip_data is None: + return jsonify({"ok": False, "error": "Not found"}), 404 + cleanup(sid) + return send_file(io.BytesIO(zip_data), mimetype="application/zip", + as_attachment=True, download_name="drhider_output.zip") diff --git a/site/routes/process_bp.py b/site/routes/process_bp.py new file mode 100644 index 0000000..06062d5 --- /dev/null +++ b/site/routes/process_bp.py @@ -0,0 +1,47 @@ +""" +POST /api/process/{sid} — обработать ВСЕ файлы сессии по одному. + +Блокирует до завершения. Результат сохраняется в сессии. +""" + +import io, zipfile, traceback +from flask import Blueprint, jsonify +from drhider import obfuscate_files, LLMClient +from drhider.builder import build_zip, build_mapping_csv +from session import get_files, store_result + +process_bp = Blueprint("process", __name__, url_prefix="/api") + + +@process_bp.route("/process/", methods=["POST"]) +def process(sid): + 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 + 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, mapping_csv=csv_str) + store_result(sid, final_zip) + return jsonify({"ok": True, "status": "done", "files": len(all_results)}) + except Exception as e: + traceback.print_exc() + return jsonify({"ok": False, "error": str(e)}), 500 diff --git a/site/routes/upload_bp.py b/site/routes/upload_bp.py new file mode 100644 index 0000000..0ac0ce6 --- /dev/null +++ b/site/routes/upload_bp.py @@ -0,0 +1,28 @@ +""" +POST /api/upload — загрузить ОДИН файл. + +Каждый вызов — новое соединение, передал → закрыл. +Сессия связывает файлы через session_id. +""" + +from flask import Blueprint, request, jsonify +from session import create_session, add_file, file_count + +upload_bp = Blueprint("upload", __name__, url_prefix="/api") + + +@upload_bp.route("/upload", methods=["POST"]) +def upload(): + 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 + return jsonify({"ok": True, "session": sid, "count": file_count(sid)}) diff --git a/tests/test_upload.py b/tests/test_upload.py new file mode 100644 index 0000000..105b785 --- /dev/null +++ b/tests/test_upload.py @@ -0,0 +1,283 @@ +""" +Тесты POST /api/upload — загрузка файлов в сессию. + +Запуск: + python3 tests/test_upload.py +""" + +import io +import os +import sys +import time +import json +import gc + +# Настраиваем пути +_site = os.path.join(os.path.dirname(__file__), "..", "site") +sys.path.insert(0, _site) +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +from flask import Flask + +# Импортируем напрямую, минуя routes/__init__.py (process_bp тянет drhider) +import importlib.util +_spec = importlib.util.spec_from_file_location( + "upload_bp", os.path.join(_site, "routes", "upload_bp.py") +) +_upload_mod = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_upload_mod) +upload_bp = _upload_mod.upload_bp + + +# ═══════════════════════════════════════════════════════════════════════════ +# Helpers +# ═══════════════════════════════════════════════════════════════════════════ + +def _new_client(): + """Создать свежий тестовый клиент.""" + app = Flask(__name__) + app.register_blueprint(upload_bp) + return app.test_client() + + +def _get_rss(): + """RSS процесса в байтах (Linux).""" + try: + with open("/proc/self/statm", "r") as f: + parts = f.read().split() + return int(parts[1]) * 4096 + except Exception: + return 0 + + +# ═══════════════════════════════════════════════════════════════════════════ +# Базовые тесты +# ═══════════════════════════════════════════════════════════════════════════ + +def test_no_file(): + """POST без файла → 400.""" + c = _new_client() + rv = c.post("/api/upload") + assert rv.status_code == 400 + data = json.loads(rv.data) + assert data["ok"] is False + assert "No file" in data["error"] + + +def test_no_filename(): + """POST с файлом без имени → 400.""" + c = _new_client() + data = {"files": (io.BytesIO(b"hello"), "")} + rv = c.post("/api/upload", data=data, content_type="multipart/form-data") + assert rv.status_code == 400 + resp = json.loads(rv.data) + assert resp["ok"] is False + assert "filename" in resp["error"].lower() + + +def test_single_file(): + """POST с одним файлом → 200, сессия + count=1.""" + c = _new_client() + data = {"files": (io.BytesIO(b"hello world"), "test.txt")} + rv = c.post("/api/upload", data=data, content_type="multipart/form-data") + assert rv.status_code == 200 + resp = json.loads(rv.data) + assert resp["ok"] is True + assert len(resp["session"]) == 12 + assert resp["count"] == 1 + + +def test_two_files_same_session(): + """Два файла в одну сессию → count=2.""" + c = _new_client() + data1 = {"files": (io.BytesIO(b"first"), "file1.txt")} + rv1 = c.post("/api/upload", data=data1, content_type="multipart/form-data") + assert rv1.status_code == 200 + sid = json.loads(rv1.data)["session"] + + data2 = { + "files": (io.BytesIO(b"second"), "file2.txt"), + "session": sid, + } + rv2 = c.post("/api/upload", data=data2, content_type="multipart/form-data") + assert rv2.status_code == 200 + resp2 = json.loads(rv2.data) + assert resp2["session"] == sid + assert resp2["count"] == 2 + + +def test_fake_session_404(): + """Несуществующая сессия → 404.""" + c = _new_client() + data = { + "files": (io.BytesIO(b"data"), "f.txt"), + "session": "deadbeef1234", + } + rv = c.post("/api/upload", data=data, content_type="multipart/form-data") + assert rv.status_code == 404 + resp = json.loads(rv.data) + assert resp["ok"] is False + + +def test_empty_file(): + """Пустой файл (0 байт) → 200.""" + c = _new_client() + data = {"files": (io.BytesIO(b""), "empty.txt")} + rv = c.post("/api/upload", data=data, content_type="multipart/form-data") + assert rv.status_code == 200 + resp = json.loads(rv.data) + assert resp["ok"] is True + assert resp["count"] == 1 + + +def test_cyrillic_filename(): + """Кириллическое имя файла → 200.""" + c = _new_client() + content = "кириллица".encode("utf-8") + data = {"files": (io.BytesIO(content), "договор_№123.txt")} + rv = c.post("/api/upload", data=data, content_type="multipart/form-data") + assert rv.status_code == 200 + resp = json.loads(rv.data) + assert resp["ok"] is True + assert resp["count"] == 1 + + +def test_content_type_json(): + """Ответ — application/json.""" + c = _new_client() + data = {"files": (io.BytesIO(b"x"), "f.txt")} + rv = c.post("/api/upload", data=data, content_type="multipart/form-data") + assert rv.status_code == 200 + assert "application/json" in rv.content_type + + +# ═══════════════════════════════════════════════════════════════════════════ +# Большие файлы +# ═══════════════════════════════════════════════════════════════════════════ + +def test_large_10mb(): + """Файл 10 MB — загружается без ошибок.""" + c = _new_client() + size = 10 * 1024 * 1024 + big_data = b"A" * size + data = {"files": (io.BytesIO(big_data), "big_10mb.bin")} + rv = c.post("/api/upload", data=data, content_type="multipart/form-data") + assert rv.status_code == 200 + resp = json.loads(rv.data) + assert resp["ok"] is True + assert resp["count"] == 1 + + +def test_large_50mb(): + """Файл 50 MB — загружается без ошибок.""" + c = _new_client() + size = 50 * 1024 * 1024 + big_data = b"B" * size + data = {"files": (io.BytesIO(big_data), "big_50mb.bin")} + rv = c.post("/api/upload", data=data, content_type="multipart/form-data") + assert rv.status_code == 200 + resp = json.loads(rv.data) + assert resp["ok"] is True + assert resp["count"] == 1 + + +# ═══════════════════════════════════════════════════════════════════════════ +# Много мелких файлов +# ═══════════════════════════════════════════════════════════════════════════ + +def test_many_small_files(): + """100 файлов по 1 KB в одну сессию.""" + c = _new_client() + data1 = {"files": (io.BytesIO(b"x" * 1024), "file_000.txt")} + rv1 = c.post("/api/upload", data=data1, content_type="multipart/form-data") + sid = json.loads(rv1.data)["session"] + + for i in range(1, 100): + data = { + "files": (io.BytesIO(b"y" * 1024), f"file_{i:03d}.txt"), + "session": sid, + } + rv = c.post("/api/upload", data=data, content_type="multipart/form-data") + assert rv.status_code == 200 + resp = json.loads(rv.data) + assert resp["session"] == sid + assert resp["count"] == i + 1 + + assert resp["count"] == 100 + + +# ═══════════════════════════════════════════════════════════════════════════ +# Утечка памяти +# ═══════════════════════════════════════════════════════════════════════════ + +def test_memory_after_large_files(): + """После загрузки больших файлов память не течёт.""" + c = _new_client() + gc.collect() + mem_before = _get_rss() + + for i in range(10): + data = {"files": (io.BytesIO(b"M" * (1024 * 1024)), f"memtest_{i}.bin")} + rv = c.post("/api/upload", data=data, content_type="multipart/form-data") + assert rv.status_code == 200 + + gc.collect() + mem_after = _get_rss() + diff_mb = (mem_after - mem_before) / (1024 * 1024) + assert diff_mb < 50, f"Memory grew {diff_mb:.1f} MB, expected < 50 MB" + + +# ═══════════════════════════════════════════════════════════════════════════ +# Бенчмарк скорости +# ═══════════════════════════════════════════════════════════════════════════ + +def test_upload_speed(): + """1 MB файл — загрузка < 500 ms.""" + c = _new_client() + data = {"files": (io.BytesIO(b"S" * (1024 * 1024)), "speed.bin")} + + start = time.time() + rv = c.post("/api/upload", data=data, content_type="multipart/form-data") + elapsed = time.time() - start + + assert rv.status_code == 200 + assert elapsed < 0.5, f"Upload took {elapsed:.3f}s, expected < 0.5s" + + +# ═══════════════════════════════════════════════════════════════════════════ +# Основной блок +# ═══════════════════════════════════════════════════════════════════════════ + +if __name__ == "__main__": + passed = 0 + failed = 0 + tests = [ + test_no_file, + test_no_filename, + test_single_file, + test_two_files_same_session, + test_fake_session_404, + test_empty_file, + test_cyrillic_filename, + test_content_type_json, + test_large_10mb, + test_large_50mb, + test_many_small_files, + test_memory_after_large_files, + test_upload_speed, + ] + for test in tests: + try: + test() + print(f" ✅ {test.__name__}") + passed += 1 + except Exception as e: + print(f" ❌ {test.__name__}: {e}") + failed += 1 + print(f"\n{'='*50}") + print(f"Результат: {passed} OK, {failed} FAIL из {len(tests)}") + if failed == 0: + print("✅ ВСЕ ТЕСТЫ ПРОЙДЕНЫ") + else: + print("❌ ЕСТЬ ПРОВАЛЫ") + sys.exit(1)