v0.0.2: split api_bp into upload/process/download modules, Connection: close, upload tests (13 new, 106 total)
Deploy drhider / validate (push) Waiting to run
Deploy drhider / validate (push) Waiting to run
This commit is contained in:
+2
-1
@@ -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
|
||||
|
||||
+6
-13
@@ -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)
|
||||
|
||||
@@ -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/<sid>", 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/<sid>", 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")
|
||||
@@ -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/<sid>", 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")
|
||||
@@ -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/<sid>", 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
|
||||
@@ -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)})
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user