MeshSate/app/main.py

"""MeshSafe 服务入口（v1.1 — 接口对齐对外文档）。

启动方式::

    uvicorn app.main:app --host 0.0.0.0 --port 8443
"""

from __future__ import annotations

import hashlib
import json
import logging
import os
import queue
import tempfile
import threading
import time
import uuid
from datetime import datetime, timezone
from typing import Any, Dict, Optional

from fastapi import Body, Depends, FastAPI, File, Form, HTTPException, UploadFile, status

from app.config import get_settings
from app.engines.structure import StructureEngine
from app.logging_config import setup_logging
from app.security.deps import (
    authenticate,
    ResponseSigningMiddleware,
    ClientContext,
)
from app.security.ratelimit import rate_limit
from app.services.policy import PolicyEngine, PolicyStore
from app.services.scanner import Scanner, ScanOutcome

# ---------------------------------------------------------------------------
# 启动时完成配置加载 + 日志初始化
# ---------------------------------------------------------------------------
setup_logging()
logger = logging.getLogger(__name__)
settings = get_settings()

# ---------------------------------------------------------------------------
# 策略存储 + 扫描器（应用级单例）
# ---------------------------------------------------------------------------
_policy_file = os.environ.get("MESHSAFE_POLICY_FILE", "config/policy.yaml")
policy_store = PolicyStore(policy_file=_policy_file)
scanner = Scanner(policy_store)

# ---------------------------------------------------------------------------
# 异步扫描队列（内存实现，生产可替换为 Redis）
# ---------------------------------------------------------------------------
_task_store: Dict[str, Dict[str, Any]] = {}
_task_lock = threading.Lock()
_task_queue: "queue.Queue[tuple]" = queue.Queue(maxsize=256)
_WORKER_COUNT = int(os.environ.get("MESHSAFE_WORKERS", "2"))


def _worker_loop() -> None:
    """后台 worker：从队列取任务执行扫描，结果存入 _task_store。"""
    local_logger = logging.getLogger("app.worker")
    while True:
        try:
            task_id, file_path, original_filename, channel, client_id = _task_queue.get()
        except Exception:
            continue

        with _task_lock:
            _task_store[task_id]["status"] = "running"

        try:
            outcome = scanner.scan_file(
                file_path=file_path,
                original_filename=original_filename,
                task_id=task_id,
                channel=channel,
                client_id=client_id,
                mode="async",
            )
            try:
                os.unlink(file_path)
            except OSError:
                pass

            with _task_lock:
                _task_store[task_id].update({
                    "status": "completed",
                    "result": outcome.to_response(include_details=True),
                    "finished_at": time.time(),
                })
            local_logger.info("异步扫描完成 task_id=%s verdict=%s action=%s",
                              task_id, outcome.verdict, outcome.action)

        except Exception as exc:  # noqa: BLE001
            local_logger.exception("异步扫描失败 task_id=%s", task_id)
            with _task_lock:
                _task_store[task_id].update({
                    "status": "failed",
                    "error": str(exc),
                    "finished_at": time.time(),
                })


def _start_workers() -> None:
    for i in range(_WORKER_COUNT):
        t = threading.Thread(target=_worker_loop, name=f"scan-worker-{i}", daemon=True)
        t.start()
    logger.info("启动 %d 个异步扫描 worker 线程", _WORKER_COUNT)


def _file_size_limit() -> int:
    """当前策略允许的最大文件大小（字节）。"""
    g = policy_store.current().raw.get("global", {})
    return int(g.get("max_file_size", 1610612736))


def _sync_max_size() -> int:
    """超过这个大小自动走异步（字节）。"""
    g = policy_store.current().raw.get("global", {})
    return int(g.get("sync_max_size", 536870912))


app = FastAPI(
    title="MeshSafe 文件安全检测服务",
    version="1.1.0",
    description=(
        "为光闸摆渡场景设计的独立文件安全检测服务。"
        '对外只做一件事：判断"这个文件能不能传进内网"，'
        "把检测结论交还给摆渡系统决定。"
    ),
)


# =====================================================================
# 全局异常处理 —— /console/* 页面 401 友好跳转登录页
# =====================================================================
from fastapi import Request
from fastapi.responses import JSONResponse, RedirectResponse
from starlette.exceptions import HTTPException as StarletteHTTPException


@app.exception_handler(HTTPException)
async def _http_exception_handler(request: Request, exc: HTTPException):
    """统一处理 HTTPException。

    对于 /console/* 路由：
      - 401 → 302 跳登录页（浏览器）或 JSON（API）
      - 403 → 提示 admin 权限不够
    对于 /api/* 路由：
      - 保持 JSON 原样返回（API 调用方需要结构化错误）
    """
    path = request.url.path

    # 控制台路由的浏览器请求 —— 401/403 友好化
    if path.startswith("/console/") and path != "/console/login":
        accept = request.headers.get("accept", "")
        want_html = "text/html" in accept or "*/*" in accept

        if exc.status_code == 401 and want_html:
            return RedirectResponse(
                url="/console/login?reason=session_expired",
                status_code=302,
            )
        if exc.status_code == 403 and want_html:
            return JSONResponse(
                status_code=403,
                content={"detail": "需要 admin 权限才能执行此操作"},
            )

    # 其他情况保持默认 JSON 返回
    return JSONResponse(status_code=exc.status_code, content={"detail": exc.detail})


# 响应签名中间件（为所有 /api/* 响应添加 HMAC 签名）
app.add_middleware(ResponseSigningMiddleware)

# 控制台路由（P2，无签名要求）
try:
    from app.console import router as console_router, diag_api_router
    app.include_router(console_router)
    app.include_router(diag_api_router)
except Exception:  # noqa: BLE001
    logger.warning("控制台路由加载失败（Jinja2 未安装或 DB 不可用）", exc_info=True)

# 静态资源（logo 等）
from fastapi.staticfiles import StaticFiles as _StaticFiles
_project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
_static_dir = os.path.join(_project_root, "static")
os.makedirs(_static_dir, exist_ok=True)
app.mount("/static", _StaticFiles(directory=_static_dir), name="static")


@app.on_event("startup")
async def _startup() -> None:
    logger.info("MeshSafe 启动中：env=%s port=%d", settings.env, settings.port)
    try:
        policy = policy_store.current()
        logger.info(
            "策略加载成功 version=%s zero_tolerance_rules=%d macro_exceptions=%d",
            policy.version,
            len(policy.zero_tolerance_rules),
            len(policy.macro_exceptions or []),
        )
    except Exception:  # noqa: BLE001
        logger.exception("策略加载失败，服务将处于 fail_closed 状态")
    _start_workers()

    # ---------- 诊断套件启动 ----------
    try:
        from app.diagnostic.config import get_diag_config
        from app.diagnostic.docker_probe import detect_docker_access
        from app.diagnostic.scheduler import start_timers

        diag_cfg = get_diag_config()
        if diag_cfg.enabled:
            docker_info = detect_docker_access(mode=diag_cfg.docker_mode)
            start_timers(docker_info["status"], docker_info)
            logger.info(
                "MeshSafe Diagnostics 已启用: docker=%s status=%s",
                diag_cfg.docker_mode, docker_info["status"],
            )
        else:
            logger.info("MeshSafe Diagnostics: disabled（config.enabled=false）")
    except Exception:  # noqa: BLE001
        logger.exception("诊断套件启动失败（不影响主服务）")


# ===========================================================================
# API 路由（对外文档 v1.1）
# ===========================================================================

# ---------------------------------------------------------------------------
# 4.5 健康检查
# ---------------------------------------------------------------------------
@app.get("/api/v1/health", tags=["health"])
async def health() -> Dict[str, Any]:
    """健康检查与引擎状态（公开接口，无需签名）。"""
    engines = scanner.health()
    engine_map: Dict[str, Any] = {}
    for name, info in engines.items():
        entry = {"available": info.get("ok", info.get("available", False))}
        if "version" in info:
            entry["version"] = info["version"]
        if "signature_date" in info:
            entry["signature_updated_at"] = info["signature_date"]
        if "rules_loaded" in info:
            entry["rules_loaded"] = info["rules_loaded"]
        if "rules_dir" in info:
            entry["rules_dir"] = info["rules_dir"]
        engine_map[name] = entry

    with _task_lock:
        pending_count = sum(1 for t in _task_store.values() if t["status"] in ("pending", "running"))

    return {
        "status": "healthy",
        "policy_version": policy_store.current().version,
        "engines": engine_map,
        "queue": {
            "pending": pending_count,
            "running": pending_count,
        },
        "server_time": datetime.now(timezone.utc).isoformat(),
    }


# ---------------------------------------------------------------------------
# 4.1 提交扫描（同步/异步自动切换）
# ---------------------------------------------------------------------------
@app.post("/api/v1/scans", tags=["scan"])
async def scans_submit(
    file: UploadFile = File(...),
    channel: str = Form("default"),
    priority: str = Form("normal"),
    ctx: ClientContext = Depends(authenticate),
    _ = Depends(rate_limit(max_requests_per_min=60, label="scans")),
) -> Dict[str, Any]:
    """提交文件进行安全检测。

    请求体相同，服务根据文件大小自动选择同步或异步：
      - ≤ sync_max_size（默认 512MB）→ 同步：请求内返回完整结论
      - > sync_max_size → 异步：返回 task_id，轮询 GET /api/v1/scans/{task_id}
    """
    client_id = ctx.client_id
    task_id = uuid.uuid4().hex

    # 1) 保存到临时文件
    suffix = os.path.splitext(file.filename or "")[1] or ""
    try:
        content = await file.read()
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
            tmp.write(content)
            tmp_path = tmp.name
    except Exception as exc:  # noqa: BLE001
        logger.exception("接收上传文件失败 task_id=%s", task_id)
        raise HTTPException(status_code=400, detail=f"文件接收失败: {exc}")

    file_size = len(content)
    del content  # 尽快释放内存

    # 2) 文件大小硬上限（超过直接 413）
    if file_size > _file_size_limit():
        try:
            os.unlink(tmp_path)
        except OSError:
            pass
        raise HTTPException(
            status_code=413,
            detail=f"文件大小 {file_size/1024/1024:.0f}MB 超过上限 {_file_size_limit()/1024/1024:.0f}MB",
        )

    # 3) 自动选择同步 / 异步
    if file_size <= _sync_max_size():
        # ---- 同步扫描 ----
        import asyncio as _asyncio
        try:
            loop = _asyncio.get_running_loop()
            outcome = await loop.run_in_executor(
                None,
                lambda: scanner.scan_file(
                    file_path=tmp_path,
                    original_filename=file.filename or f"upload{suffix}",
                    task_id=task_id,
                    channel=channel,
                    client_id=client_id,
                ),
            )
        finally:
            try:
                os.unlink(tmp_path)
            except OSError:
                pass

        # 同步结果也存入 _task_store，支持后续 GET /api/v1/scans/{task_id} 查询和 review
        with _task_lock:
            _task_store[task_id] = {
                "status": "completed",
                "result": outcome.to_response(include_details=True),
                "finished_at": time.time(),
                "filename": file.filename,
                "channel": channel,
            }

        return outcome.to_response(include_details=True)

    else:
        # ---- 异步扫描 ----
        with _task_lock:
            if len(_task_store) >= 1024:
                finished = [tid for tid, t in _task_store.items()
                            if t["status"] in ("completed", "failed")]
                for tid in finished[-512:]:
                    del _task_store[tid]

            _task_store[task_id] = {
                "status": "pending",
                "submitted_at": time.time(),
                "filename": file.filename,
                "channel": channel,
            }

        try:
            _task_queue.put_nowait((task_id, tmp_path, file.filename or f"upload{suffix}",
                                     channel, client_id))
        except queue.Full:
            with _task_lock:
                _task_store[task_id] = {"status": "failed", "error": "扫描队列已满"}
            try:
                os.unlink(tmp_path)
            except OSError:
                pass
            raise HTTPException(status_code=503, detail="扫描队列已满，请稍后重试")

        return {
            "task_id": task_id,
            "status": "pending",
            "message": "任务已受理，请轮询 /api/v1/scans/{task_id} 获取结果",
        }


# ---------------------------------------------------------------------------
# 4.2 查询扫描结果
# ---------------------------------------------------------------------------
@app.get("/api/v1/scans/{task_id}", tags=["scan"])
async def scans_get(task_id: str, ctx: ClientContext = Depends(authenticate)) -> Dict[str, Any]:
    """查询扫描结果。同步模式下用于重试补查，异步模式下用于轮询。"""
    with _task_lock:
        task = _task_store.get(task_id)

    if not task:
        raise HTTPException(status_code=404, detail=f"任务 {task_id} 不存在或已过期")

    status_val = task["status"]

    if status_val == "completed":
        return task["result"]

    if status_val == "failed":
        return {
            "task_id": task_id,
            "status": "failed",
            "verdict": "error",
            "action": "error",
            "score": 0.0,
            "reason": task.get("error", "未知错误"),
            "error": task.get("error"),
            "finished_at": datetime.fromtimestamp(task.get("finished_at", time.time()), tz=timezone.utc).isoformat(),
        }

    # pending / running → 返回占位结构
    return {
        "task_id": task_id,
        "status": status_val,
        "verdict": "pending",
        "action": "pending",
        "score": 0.0,
        "reason": "",
    }


# ---------------------------------------------------------------------------
# 4.3 重新扫描
# ---------------------------------------------------------------------------
@app.post("/api/v1/scans/{task_id}/rescan", tags=["scan"])
async def scans_rescan(
    task_id: str,
    file: UploadFile = File(...),
    ctx: ClientContext = Depends(authenticate),
) -> Dict[str, Any]:
    """使用当前生效策略重新扫描。

    原始文件可能已被清理（进程退出即 unlink），因此需要重新上传文件。
    """
    with _task_lock:
        old_task = _task_store.get(task_id)

    suffix = os.path.splitext(file.filename or "")[1] or ""
    try:
        content = await file.read()
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
            tmp.write(content)
            tmp_path = tmp.name
    except Exception as exc:  # noqa: BLE001
        logger.exception("rescan 接收失败 task_id=%s", task_id)
        raise HTTPException(status_code=400, detail=f"文件接收失败: {exc}")

    try:
        import asyncio as _asyncio
        loop = _asyncio.get_running_loop()
        outcome = await loop.run_in_executor(
            None,
            lambda: scanner.scan_file(
                file_path=tmp_path,
                original_filename=file.filename or (old_task.get("filename", f"upload{suffix}") if old_task else f"upload{suffix}"),
                task_id=task_id,
                channel=old_task.get("channel", "default") if old_task else "default",
                client_id=ctx.client_id,
            ),
        )
    finally:
        try:
            os.unlink(tmp_path)
        except OSError:
            pass

    return outcome.to_response(include_details=True)


# ---------------------------------------------------------------------------
# 4.4 人工复核回写
# ---------------------------------------------------------------------------
@app.post("/api/v1/scans/{task_id}/review", tags=["scan"])
async def scans_review(
    task_id: str,
    body: Dict[str, Any] = Body(...),
    ctx: ClientContext = Depends(authenticate),
) -> Dict[str, Any]:
    """回写人工复核结论。

    请求体::

        { "decision": "pass" | "reject", "comment": "..." }
    """
    decision = body.get("decision", "").lower()
    comment = body.get("comment", "")

    if decision not in ("pass", "reject"):
        raise HTTPException(status_code=400, detail="decision 必须是 'pass' 或 'reject'")
    if not comment.strip():
        raise HTTPException(status_code=400, detail="comment 必填（审计要求）")

    with _task_lock:
        task = _task_store.get(task_id)
        if task:
            task["review_decision"] = decision
            task["review_comment"] = comment
            task["reviewed_by"] = ctx.client_id
            task["reviewed_at"] = time.time()

    return {
        "task_id": task_id,
        "decision": decision,
        "comment": comment,
        "recorded_at": datetime.now(timezone.utc).isoformat(),
        "reviewed_by": ctx.client_id,
    }


# ---------------------------------------------------------------------------
# 4.6 查询当前策略
# ---------------------------------------------------------------------------
@app.get("/api/v1/policies", tags=["policy"])
async def policies_get(ctx: ClientContext = Depends(authenticate)) -> Dict[str, Any]:
    """查询当前生效的全局策略（只读）。"""
    policy = policy_store.current()
    return {
        "version": policy.version,
        "sha256": policy.sha256,
        "content": policy.raw,
    }


# ---------------------------------------------------------------------------
# 4.7 策略热重载
# ---------------------------------------------------------------------------
@app.post("/api/v1/policies/reload", tags=["policy"])
async def policies_reload(ctx: ClientContext = Depends(authenticate)) -> Dict[str, Any]:
    """从策略文件重新加载（无需重启）。"""
    try:
        policy = policy_store.reload()
        logger.info("策略热重载成功 version=%s", policy.version)
        return {
            "reloaded": True,
            "version": policy.version,
            "sha256": policy.sha256,
        }
    except Exception as exc:  # noqa: BLE001
        logger.exception("策略热重载失败")
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"策略加载失败: {exc}",
        )


# ---------------------------------------------------------------------------
# 4.8 审计日志完整性校验
# ---------------------------------------------------------------------------
@app.get("/api/v1/audit/verify", tags=["audit"])
async def audit_verify(ctx: ClientContext = Depends(authenticate)) -> Dict[str, Any]:
    """校验审计日志哈希链完整性。

    当前实现为 stub（哈希链待二期接入）—— 返回 intact=True 表示审计链路完好。
    """
    return {
        "checked": 0,
        "intact": True,
        "broken_at_id": None,
        "verified_at": datetime.now(timezone.utc).isoformat(),
        "note": "审计哈希链为二期功能，当前返回 intact=True 占位",
    }


# ---------------------------------------------------------------------------
# 直接运行支持
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app.main:app", host=settings.host, port=settings.port, reload=False)