aps-agent/server/auth/middleware.py

from __future__ import annotations

import hashlib
import os
import re

from fastapi import Request
from fastapi.responses import JSONResponse
from starlette.datastructures import MutableHeaders

from server.auth import agent_tokens
from server.auth.context import IdentityContext, bind_identity, reset_identity
from server.auth.licenses import get_license_provider, is_desktop_request
from server.auth.providers import AuthError, get_auth_provider

PUBLIC_API_PATHS = {
    "/api/health",
    "/api/auth/login",
    "/api/auth/login/precheck",
    "/api/auth/tenants",
    "/api/auth/captcha",
    "/api/auth/refresh",
    "/api/auth/logout",
    "/api/auth/license/activate",
}

SELF_AUTHORIZED_WRITE_PREFIXES = (
    "/api/workspace",
    "/api/projects",
    "/api/sessions",
    "/api/project-files",
    "/api/drawings",  # 图纸解析只读放行；候选入库走 P2 门禁
)

ANONYMOUS_COOKIE = "aps_anonymous_id"
ANONYMOUS_COOKIE_MAX_AGE = 365 * 24 * 3600


def authentication_enabled() -> bool:
    """Return whether browser/desktop requests must present login credentials."""
    value = (os.environ.get("APS_AUTH_ENABLED") or "1").strip().lower()
    return value not in {"0", "false", "no", "off"}


def bypass_identity() -> IdentityContext:
    """Single local administrator identity used only when auth is explicitly disabled."""
    return IdentityContext(
        user_id=1,
        username="local-admin",
        fullname="本地免登录管理员",
        tenant_uuid="platform",
        roles=(
            "system",
            "admin",
            "planner",
            "approver",
            "auditor",
            "scheduler",
            "desktop",
        ),
        auth_kind="disabled",
    )


def _visitor_identity(visitor: str) -> IdentityContext | None:
    visitor = str(visitor or "").strip().lower()
    if (
        len(visitor) < 8
        or len(visitor) > 128
        or not re.fullmatch(r"[a-z0-9_-]+", visitor)
    ):
        return None
    digest = hashlib.sha256(f"aps-anonymous:{visitor}".encode()).hexdigest()
    user_id = int(digest[:15], 16)
    return IdentityContext(
        user_id=user_id,
        username=f"anon-{digest[:10]}",
        fullname=f"访客 {digest[:8]}",
        tenant_uuid="platform",
        roles=(
            "system",
            "admin",
            "planner",
            "approver",
            "auditor",
            "scheduler",
            "desktop",
        ),
        auth_kind="anonymous",
    )


def anonymous_visitor_identity(request: Request) -> tuple[IdentityContext, str | None]:
    """免登录云部署：按浏览器访客 ID 派生独立用户作用域。

    Web 前端把 localStorage 里稳定的访客 ID 随 X-APS-Visitor-ID 发送；
    后端用同一 ID 派生出固定 user_id，并签发长期 HttpOnly Cookie。
    之后即使 localStorage 被清空，Cookie 仍能让浏览器找回同一身份。
    没有 Cookie/头的旧客户端回退到本地管理员身份，保持既有行为。
    """
    identity = _visitor_identity(request.cookies.get(ANONYMOUS_COOKIE))
    if identity is not None:
        return identity, None
    visitor = request.headers.get("x-aps-visitor-id")
    identity = _visitor_identity(visitor)
    if identity is not None:
        return identity, str(visitor).strip().lower()
    return bypass_identity(), None


class AuthenticationMiddleware:
    """ASGI middleware keeps identity bound for the full streaming response lifetime."""

    def __init__(self, app) -> None:
        self.app = app

    async def __call__(self, scope, receive, send) -> None:
        if scope.get("type") != "http":
            await self.app(scope, receive, send)
            return
        path = scope.get("path") or ""
        if (
            scope.get("method") == "OPTIONS"
            or not path.startswith("/api/")
            or path in PUBLIC_API_PATHS
        ):
            await self.app(scope, receive, send)
            return
        request = Request(scope, receive=receive)
        anon_cookie: str | None = None
        authorization = request.headers.get("authorization") or ""
        agent_request = path.startswith("/api/agent/")
        if (
            agent_request
            and not authorization.lower().startswith("bearer ")
            and not authentication_enabled()
        ):
            response = JSONResponse(
                {
                    "error": {
                        "code": "AGENT_TOKEN_REQUIRED",
                        "message": "Pi Agent 端点需要 Authorization: Bearer <agent-token>",
                    }
                },
                status_code=401,
            )
            await response(scope, receive, send)
            return
        if agent_request and authorization.lower().startswith("bearer "):
            try:
                identity = agent_tokens.authenticate(authorization[7:].strip())
            except AuthError as exc:
                response = JSONResponse(
                    {"error": {"code": exc.code, "message": exc.message}},
                    status_code=exc.status_code,
                )
                await response(scope, receive, send)
                return
        elif authentication_enabled():
            try:
                provider = (
                    get_license_provider()
                    if is_desktop_request(request)
                    else get_auth_provider()
                )
                identity = await provider.authenticate(request)
            except AuthError as exc:
                response = JSONResponse(
                    {"error": {"code": exc.code, "message": exc.message}},
                    status_code=exc.status_code,
                )
                await response(scope, receive, send)
                return
        else:
            identity, anon_cookie = anonymous_visitor_identity(request)
        token = bind_identity(identity)
        try:
            method = scope.get("method") or "GET"
            if method not in {"GET", "HEAD", "OPTIONS"} and not path.startswith(
                SELF_AUTHORIZED_WRITE_PREFIXES
            ):
                from server.state.projects import get_project_store

                try:
                    get_project_store().require_active_write()
                except PermissionError as exc:
                    response = JSONResponse(
                        {"error": {"code": "PROJECT_READ_ONLY", "message": str(exc)}},
                        status_code=403,
                    )
                    await response(scope, receive, send)
                    return
            if anon_cookie:
                secure = "; Secure" if request.url.scheme == "https" else ""
                cookie_value = (
                    f"{ANONYMOUS_COOKIE}={anon_cookie}; Path=/; "
                    f"Max-Age={ANONYMOUS_COOKIE_MAX_AGE}; HttpOnly; SameSite=Lax{secure}"
                )

                async def send_with_cookie(message) -> None:
                    if message["type"] == "http.response.start":
                        MutableHeaders(scope=message).append("set-cookie", cookie_value)
                    await send(message)

                await self.app(scope, receive, send_with_cookie)
            else:
                await self.app(scope, receive, send)
        finally:
            reset_identity(token)