feat: интегрировать модуль upload в drhider (бэк blueprint + фронт ES-модули)
This commit is contained in:
“Naeel”
+6
-113
@@ -17,41 +17,19 @@ import time
|
||||
import zipfile
|
||||
import traceback
|
||||
import logging
|
||||
import httpx
|
||||
from datetime import datetime, timedelta
|
||||
from flask import Blueprint, request, send_file, jsonify, Response, stream_with_context
|
||||
|
||||
from drhider import obfuscate_files, LLMClient
|
||||
from session import (create_session, add_file, get_files, store_result,
|
||||
get_result, store_csv, get_csv, cleanup, file_count,
|
||||
MAX_FILE_BYTES, pause_ttl, resume_ttl,
|
||||
request_cancel, get_cancel_event)
|
||||
from upload.backend.session import (create_session, add_file, get_files, store_result,
|
||||
get_result, store_csv, get_csv, cleanup, file_count,
|
||||
MAX_FILE_BYTES, pause_ttl, resume_ttl,
|
||||
request_cancel, get_cancel_event)
|
||||
from upload.backend.upload_refs import safe_name
|
||||
|
||||
api_bp = Blueprint("api", __name__, url_prefix="/api")
|
||||
log = logging.getLogger("routes.api_bp")
|
||||
|
||||
# Ретраи pull из ВМ-буфера: защита от разовых DNS/сетевых сбоев (gaierror -5 и т.п.)
|
||||
PULL_RETRIES = 3
|
||||
PULL_RETRY_DELAY = 2 # секунды между попытками
|
||||
|
||||
# Доверенный префикс ВМ-буфера — валидация URL при pull (защита от SSRF)
|
||||
VM_UPLOAD_PREFIX = "https://contracts.kube5s.ru/drhider-upload/"
|
||||
|
||||
|
||||
def _safe_name(name: str) -> str:
|
||||
"""Санитизировать имя файла: защита от path traversal, сохраняя подпапки.
|
||||
|
||||
Запрещает '..' и абсолютные пути; нормализует слэши. Возвращает "" если
|
||||
имя пустое или небезопасное.
|
||||
"""
|
||||
if not name:
|
||||
return ""
|
||||
name = name.replace("\\", "/")
|
||||
parts = [p for p in name.split("/") if p and p != "."]
|
||||
if not parts or any(p == ".." for p in parts):
|
||||
return ""
|
||||
return "/".join(parts)
|
||||
|
||||
|
||||
def _disconnect_exceptions():
|
||||
"""Исключения, означающие отключение клиента SSE."""
|
||||
@@ -72,7 +50,7 @@ def upload():
|
||||
added = 0
|
||||
had_unnamed = False
|
||||
for f in uploaded:
|
||||
name = _safe_name(f.filename)
|
||||
name = safe_name(f.filename)
|
||||
if not name:
|
||||
had_unnamed = True
|
||||
continue
|
||||
@@ -96,91 +74,6 @@ def upload():
|
||||
return jsonify({"ok": True, "session": sid, "count": file_count(sid)})
|
||||
|
||||
|
||||
@api_bp.route("/upload_refs", methods=["POST"])
|
||||
def upload_refs():
|
||||
"""Принять ссылки на файлы (загружены на ВМ-буфер), забрать по egress.
|
||||
|
||||
Вход: JSON {"session": "...", "files": [{"name": str, "size": int, "url": str}]}.
|
||||
Каждый файл тянется ИСХОДЯЩИМ GET'ом с ВМ (egress не ограничен шлюзом),
|
||||
читается по частям (stream), кладётся в сессию. После успешного pull файл
|
||||
удаляется с ВМ (best-effort; TTL-чистка на ВМ тоже есть).
|
||||
"""
|
||||
data = request.get_json(silent=True) or {}
|
||||
sid = data.get("session") or create_session()
|
||||
refs = data.get("files") or []
|
||||
if not refs:
|
||||
log.warning("upload_refs: no files, sid=%s", sid)
|
||||
return jsonify({"ok": False, "error": "No files"}), 400
|
||||
added = 0
|
||||
try:
|
||||
with httpx.Client(timeout=120, follow_redirects=True) as client:
|
||||
for ref in refs:
|
||||
name = _safe_name(ref.get("name") or "")
|
||||
url = ref.get("url")
|
||||
if not name or not url:
|
||||
continue
|
||||
# SSRF-защита: тянуть можно ТОЛЬКО с доверенного ВМ-буфера
|
||||
if not url.startswith(VM_UPLOAD_PREFIX):
|
||||
log.warning("upload_refs: unsafe URL, skip sid=%s url=%r", sid, url)
|
||||
continue
|
||||
# Лимит на один файл (50 МБ): сверх лимита — пропускаем (не участвует)
|
||||
if (ref.get("size") or 0) > MAX_FILE_BYTES:
|
||||
log.warning("upload_refs: file exceeds %dMB, skip sid=%s file=%r size=%s",
|
||||
MAX_FILE_BYTES // (1024 * 1024), sid, name, ref.get("size"))
|
||||
try:
|
||||
client.delete(url)
|
||||
except Exception:
|
||||
pass
|
||||
continue
|
||||
# Pull с ретраями: разовые DNS/сетевые сбои не роняют всю загрузку
|
||||
content = None
|
||||
last_err = None
|
||||
for attempt in range(PULL_RETRIES):
|
||||
try:
|
||||
with client.stream("GET", url) as resp:
|
||||
resp.raise_for_status()
|
||||
content = b"".join(resp.iter_bytes())
|
||||
last_err = None
|
||||
break
|
||||
except Exception as e:
|
||||
last_err = e
|
||||
log.warning("upload_refs: pull attempt %d/%d failed sid=%s file=%r: %r",
|
||||
attempt + 1, PULL_RETRIES, sid, name, e)
|
||||
time.sleep(PULL_RETRY_DELAY)
|
||||
if content is None:
|
||||
raise last_err if last_err else RuntimeError("pull failed")
|
||||
log.info("upload_refs: pulled sid=%s file=%r size=%d", sid, name, len(content))
|
||||
if len(content) > MAX_FILE_BYTES:
|
||||
log.warning("upload_refs: pulled file exceeds %dMB, skip sid=%s file=%r size=%d",
|
||||
MAX_FILE_BYTES // (1024 * 1024), sid, name, len(content))
|
||||
try:
|
||||
client.delete(url)
|
||||
except Exception:
|
||||
pass
|
||||
continue
|
||||
if not add_file(sid, name, content):
|
||||
# Различить: сессия исчезла vs превышен суммарный лимит сессии
|
||||
if get_files(sid) is None:
|
||||
log.warning("upload_refs: session not found, sid=%s file=%r", sid, name)
|
||||
return jsonify({"ok": False, "error": "Session not found"}), 404
|
||||
log.warning("upload_refs: session limit exceeded, skip sid=%s file=%r", sid, name)
|
||||
try:
|
||||
client.delete(url)
|
||||
except Exception:
|
||||
pass
|
||||
continue
|
||||
try:
|
||||
client.delete(url) # убрать файл с ВМ после загрузки
|
||||
except Exception:
|
||||
pass
|
||||
added += 1
|
||||
except Exception as e:
|
||||
log.error("upload_refs: pull error sid=%s: %r", sid, e)
|
||||
return jsonify({"ok": False, "error": "Pull failed: %s" % e}), 502
|
||||
log.info("upload_refs: done sid=%s added=%d total=%d", sid, added, file_count(sid))
|
||||
return jsonify({"ok": True, "session": sid, "count": file_count(sid)})
|
||||
|
||||
|
||||
@api_bp.route("/session_files/<sid>", methods=["GET"])
|
||||
def session_files(sid):
|
||||
"""Отладка: список файлов сессии с размерами (для теста загрузки)."""
|
||||
|
||||
Reference in New Issue
Block a user