124 lines
6.6 KiB
Python
124 lines
6.6 KiB
Python
"""Переиспользуемый Blueprint слоя 2: POST /upload_refs (pull с ВМ-буфера в сессию).
|
||
|
||
Поведение 1:1 с drhider v0.0.75 (site/routes/api_bp.py):
|
||
- _safe_name (path traversal)
|
||
- SSRF-валидация url.startswith(VM_UPLOAD_PREFIX)
|
||
- лимит на один файл -> delete+skip
|
||
- pull с ретраями
|
||
- лимит сессии -> skip; отсутствие сессии -> 404
|
||
- delete url с ВМ (best-effort)
|
||
"""
|
||
|
||
import httpx
|
||
import logging
|
||
from flask import Blueprint, request, jsonify
|
||
|
||
from ..session import (create_session, add_file, get_files, file_count,
|
||
MAX_FILE_BYTES, configure)
|
||
from .config import PULL_RETRIES, PULL_RETRY_DELAY, VM_UPLOAD_PREFIX
|
||
from .safe_name import safe_name
|
||
from .pull_file import pull_file
|
||
|
||
log = logging.getLogger("upload.upload_refs")
|
||
|
||
|
||
def create_upload_refs_blueprint(cfg: dict) -> Blueprint:
|
||
"""Создать Blueprint с эндпоинтом upload_refs.
|
||
|
||
cfg (все ключи опциональны, есть дефолты):
|
||
apiPrefix (str) — префикс Blueprint, по умолчанию "/api"
|
||
vmUploadPrefix (str) — доверенный префикс ВМ-буфера (SSRF-валидация)
|
||
maxFileBytes (int) — лимит на один файл
|
||
maxSessionBytes (int) — суммарный лимит сессии (применяется к сессиям)
|
||
ttlSeconds (int) — TTL сессии
|
||
pullRetries (int) — ретраи pull
|
||
pullRetryDelay (int) — пауза между ретраями (сек)
|
||
pullTimeout (int) — таймаут одного GET pull
|
||
"""
|
||
prefix = cfg.get("apiPrefix", "/api")
|
||
vm_prefix = cfg.get("vmUploadPrefix", VM_UPLOAD_PREFIX)
|
||
max_file_bytes = cfg.get("maxFileBytes", MAX_FILE_BYTES)
|
||
pull_retries = cfg.get("pullRetries", PULL_RETRIES)
|
||
pull_delay = cfg.get("pullRetryDelay", PULL_RETRY_DELAY)
|
||
pull_timeout = cfg.get("pullTimeout", 120)
|
||
|
||
# Применить лимиты сессии/TTL из конфига (глобально для всех сессий)
|
||
configure(
|
||
max_file_bytes=cfg.get("maxFileBytes"),
|
||
max_session_bytes=cfg.get("maxSessionBytes"),
|
||
ttl_seconds=cfg.get("ttlSeconds"),
|
||
)
|
||
|
||
bp = Blueprint("upload_refs", __name__, url_prefix=prefix)
|
||
|
||
@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=pull_timeout, 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_prefix):
|
||
log.warning("upload_refs: unsafe URL, skip sid=%s url=%r", sid, url)
|
||
continue
|
||
# Лимит на один файл: сверх лимита — пропускаем (не участвует)
|
||
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 = pull_file(client, url, pull_retries, pull_delay, sid=sid, name=name) # бросает при неудаче всех попыток
|
||
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)})
|
||
|
||
return bp
|