ForcePilot/backend/package/yuxi/channel/gateway/rbac.py

import logging
from collections.abc import Callable
from enum import StrEnum
from functools import wraps
from typing import Any

logger = logging.getLogger(__name__)


class GatewayRole(StrEnum):
    SUPERADMIN = "superadmin"
    ADMIN = "admin"
    OPERATOR = "operator"
    VIEWER = "viewer"


ROLE_HIERARCHY: dict[GatewayRole, int] = {
    GatewayRole.SUPERADMIN: 4,
    GatewayRole.ADMIN: 3,
    GatewayRole.OPERATOR: 2,
    GatewayRole.VIEWER: 1,
}


_USER_ROLE_MAP: dict[str, GatewayRole] = {
    "superadmin": GatewayRole.SUPERADMIN,
    "admin": GatewayRole.ADMIN,
    "user": GatewayRole.OPERATOR,
}


def map_user_role(user_role: str | None) -> GatewayRole:
    if not user_role:
        return GatewayRole.VIEWER
    return _USER_ROLE_MAP.get(user_role.lower(), GatewayRole.VIEWER)


READ_ONLY_METHODS: set[str] = {
    "system.health",
    "system.version",
    "channels.list",
    "channels.status",
    "channels.probe",
    "chat.history",
    "sessions.list",
    "sessions.history",
    "plugins.list",
    "agentTools.list",
    "config.get",
}

OPERATOR_METHODS: set[str] = READ_ONLY_METHODS | {
    "chat.send",
    "chat.stream",
    "chat.cancel",
    "sessions.create",
    "sessions.delete",
    "allowlist.get",
    "cron.list",
    "agentTools.execute",
    "message.send",
    "message.action",
}

ADMIN_METHODS: set[str] = OPERATOR_METHODS | {
    "system.log.tail",
    "system.config",
    "channels.start",
    "channels.stop",
    "channels.restart",
    "channels.configure",
    "pairing.list",
    "pairing.approve",
    "pairing.reject",
    "allowlist.add",
    "allowlist.remove",
    "identity_links.list",
    "identity_links.add",
    "identity_links.remove",
    "plugins.install",
    "config.set",
}

SUPERADMIN_METHODS: set[str] = ADMIN_METHODS | {
    "cron.create",
    "cron.delete",
    "cron.force_run",
    "cron.pause",
    "cron.resume",
}

METHOD_MIN_ROLE: dict[str, GatewayRole] = {}

for _method in READ_ONLY_METHODS:
    METHOD_MIN_ROLE[_method] = GatewayRole.VIEWER

for _method in OPERATOR_METHODS - READ_ONLY_METHODS:
    METHOD_MIN_ROLE[_method] = GatewayRole.OPERATOR

for _method in ADMIN_METHODS - OPERATOR_METHODS:
    METHOD_MIN_ROLE[_method] = GatewayRole.ADMIN

for _method in {"cron.create", "cron.delete", "cron.force_run", "cron.pause", "cron.resume"}:
    METHOD_MIN_ROLE[_method] = GatewayRole.SUPERADMIN


def check_permission(role: GatewayRole | None, method: str) -> bool:
    effective_role = role or GatewayRole.VIEWER
    min_role = METHOD_MIN_ROLE.get(method)
    if min_role is None:
        return True
    return ROLE_HIERARCHY[effective_role] >= ROLE_HIERARCHY[min_role]


def require_role(min_role: GatewayRole):
    def decorator(
        func: Callable[..., Any],
    ) -> Callable[..., Any]:
        @wraps(func)
        async def wrapper(*args: Any, **kwargs: Any) -> Any:
            caller_role: GatewayRole | None = kwargs.pop("_caller_role", None)
            if caller_role is None:
                logger.warning(
                    "RBAC: _caller_role not found in handler kwargs for %s, allowing by default",
                    func.__name__,
                )
                return await func(*args, **kwargs)

            if ROLE_HIERARCHY[caller_role] < ROLE_HIERARCHY[min_role]:
                from yuxi.channel.gateway.protocol import GatewayErrorCode, RpcResponse

                logger.warning(
                    "RBAC: %s requires %s, caller has %s — denied",
                    func.__name__,
                    min_role.value,
                    caller_role.value,
                )
                return RpcResponse(
                    id=kwargs.get("request_id", ""),
                    ok=False,
                    error_code=GatewayErrorCode.PERMISSION_DENIED,
                    error_message=f"需要 {min_role.value} 权限，当前角色为 {caller_role.value}",
                )

            return await func(*args, **kwargs)

        return wrapper

    return decorator