contracts-flask/upload/backend/upload_refs/blueprint.py

"""Переиспользуемый 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, sink=None) -> 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

    sink (callable | None): если задан — вместо add_file в in-memory сессию
        вызывается sink(name, content, **extra) на каждый вытянутый файл,
        endpoint возвращает {"ok": true, "results": [...]}. extra — сквозные
        поля тела запроса (batch_id, contract_id, zip_source). Если None —
        прежнее поведение drhider (in-memory сессия).
    """
    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

        extra = {k: data[k] for k in ("batch_id", "contract_id", "zip_source") if data.get(k) is not None}
        results = [] if sink else None
        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)
                        if sink:
                            results.append({"name": name, "ok": False, "error": "invalid url (SSRF guard)"})
                        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
                        if sink:
                            results.append({"name": name, "ok": False, "error": "file too large"})
                        continue
                    # Pull с ретраями: разовые DNS/сетевые сбои не роняют всю загрузку
                    try:
                        content = pull_file(client, url, pull_retries, pull_delay, sid=sid, name=name)
                    except Exception as e:
                        log.warning("upload_refs: pull failed sid=%s file=%r: %r", sid, name, e)
                        if sink:
                            results.append({"name": name, "ok": False, "error": "pull failed: %s" % e})
                        continue
                    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
                        if sink:
                            results.append({"name": name, "ok": False, "error": "file too large"})
                        continue
                    if sink:
                        # Отдать файл приложению через sink (вместо in-memory сессии)
                        try:
                            client.delete(url)
                        except Exception:
                            pass
                        try:
                            r = sink(name, content, **extra) or {}
                            results.append({"name": name, **r})
                        except Exception as e:
                            log.error("upload_refs: sink failed file=%r: %r", name, e)
                            results.append({"name": name, "ok": False, "error": str(e)})
                        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

        if sink:
            return jsonify({"ok": True, "results": results})
        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