LangBot/src/langbot/pkg/api/http/controller/group.py

from __future__ import annotations

import abc
import typing
import enum
import quart
import traceback
from quart.typing import RouteCallable

from ....core import app

# Maximum file upload size limit (10MB)
MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB


preregistered_groups: list[type[RouterGroup]] = []
"""Pre-registered list of RouterGroup"""


def group_class(name: str, path: str) -> typing.Callable[[typing.Type[RouterGroup]], typing.Type[RouterGroup]]:
    """注册一个 RouterGroup"""

    def decorator(cls: typing.Type[RouterGroup]) -> typing.Type[RouterGroup]:
        cls.name = name
        cls.path = path
        preregistered_groups.append(cls)
        return cls

    return decorator


class AuthType(enum.Enum):
    """Authentication type"""

    NONE = 'none'
    USER_TOKEN = 'user-token'
    API_KEY = 'api-key'
    USER_TOKEN_OR_API_KEY = 'user-token-or-api-key'


class RouterGroup(abc.ABC):
    name: str

    path: str

    ap: app.Application

    quart_app: quart.Quart

    def __init__(self, ap: app.Application, quart_app: quart.Quart) -> None:
        self.ap = ap
        self.quart_app = quart_app

    @abc.abstractmethod
    async def initialize(self) -> None:
        pass

    def route(
        self,
        rule: str,
        auth_type: AuthType = AuthType.USER_TOKEN,
        **options: typing.Any,
    ) -> typing.Callable[[RouteCallable], RouteCallable]:  # decorator
        """Register a route"""

        def decorator(f: RouteCallable) -> RouteCallable:
            nonlocal rule
            rule = self.path + rule

            async def handler_error(*args, **kwargs):
                if auth_type == AuthType.USER_TOKEN:
                    # get token from Authorization header
                    token = quart.request.headers.get('Authorization', '').replace('Bearer ', '')

                    if not token:
                        return self.http_status(401, -1, 'No valid user token provided')

                    try:
                        user_email = await self.ap.user_service.verify_jwt_token(token)

                        # check if this account exists
                        user = await self.ap.user_service.get_user_by_email(user_email)
                        if not user:
                            return self.http_status(401, -1, 'User not found')

                        # check if f accepts user_email parameter
                        if 'user_email' in f.__code__.co_varnames:
                            kwargs['user_email'] = user_email
                    except Exception as e:
                        return self.http_status(401, -1, str(e))

                elif auth_type == AuthType.API_KEY:
                    # get API key from Authorization header or X-API-Key header
                    api_key = quart.request.headers.get('X-API-Key', '')
                    if not api_key:
                        auth_header = quart.request.headers.get('Authorization', '')
                        if auth_header.startswith('Bearer '):
                            api_key = auth_header.replace('Bearer ', '')

                    if not api_key:
                        return self.http_status(401, -1, 'No valid API key provided')

                    try:
                        is_valid = await self.ap.apikey_service.verify_api_key(api_key)
                        if not is_valid:
                            return self.http_status(401, -1, 'Invalid API key')
                    except Exception as e:
                        return self.http_status(401, -1, str(e))

                elif auth_type == AuthType.USER_TOKEN_OR_API_KEY:
                    # Try API key first (check X-API-Key header)
                    api_key = quart.request.headers.get('X-API-Key', '')

                    if api_key:
                        # API key authentication
                        try:
                            is_valid = await self.ap.apikey_service.verify_api_key(api_key)
                            if not is_valid:
                                return self.http_status(401, -1, 'Invalid API key')
                        except Exception as e:
                            return self.http_status(401, -1, str(e))
                    else:
                        # Try user token authentication (Authorization header)
                        token = quart.request.headers.get('Authorization', '').replace('Bearer ', '')

                        if not token:
                            return self.http_status(
                                401, -1, 'No valid authentication provided (user token or API key required)'
                            )

                        try:
                            user_email = await self.ap.user_service.verify_jwt_token(token)

                            # check if this account exists
                            user = await self.ap.user_service.get_user_by_email(user_email)
                            if not user:
                                return self.http_status(401, -1, 'User not found')

                            # check if f accepts user_email parameter
                            if 'user_email' in f.__code__.co_varnames:
                                kwargs['user_email'] = user_email
                        except Exception:
                            # If user token fails, maybe it's an API key in Authorization header
                            try:
                                is_valid = await self.ap.apikey_service.verify_api_key(token)
                                if not is_valid:
                                    return self.http_status(401, -1, 'Invalid authentication credentials')
                            except Exception as e:
                                return self.http_status(401, -1, str(e))

                try:
                    return await f(*args, **kwargs)

                except Exception as e:  # 自动 500
                    traceback.print_exc()
                    # return self.http_status(500, -2, str(e))
                    return self.http_status(500, -2, str(e))

            new_f = handler_error
            new_f.__name__ = (self.name + rule).replace('/', '__')
            new_f.__doc__ = f.__doc__

            self.quart_app.route(rule, **options)(new_f)
            return f

        return decorator

    def success(self, data: typing.Any = None) -> quart.Response:
        """Return a 200 response"""
        return quart.jsonify(
            {
                'code': 0,
                'msg': 'ok',
                'data': data,
            }
        )

    def fail(self, code: int, msg: str) -> quart.Response:
        """Return an error response"""

        return quart.jsonify(
            {
                'code': code,
                'msg': msg,
            }
        )

    def http_status(self, status: int, code: int, msg: str) -> typing.Tuple[quart.Response, int]:
        """返回一个指定状态码的响应"""
        return (self.fail(code, msg), status)