utils/server_util.py

import asyncio
import ipaddress
import os
import socket
import uuid
from collections.abc import Callable

import psutil
from fastapi import FastAPI, Request, applications
from fastapi.openapi.docs import get_redoc_html, get_swagger_ui_html, get_swagger_ui_oauth2_redirect_html
from fastapi.openapi.utils import get_openapi
from fastapi.responses import HTMLResponse, JSONResponse
from redis import asyncio as aioredis

from config.env import AppConfig


class APIDocsUtil:
    """
    API文档工具类
    """

    # API文档URLs
    _OPENAPI_URL = '/openapi.json'
    _PROXY_OPENAPI_URL = '/proxy-openapi.json'
    _DOCS_URL = '/docs'
    _PROXY_DOCS_URL = '/proxy-docs'
    _REDOC_URL = '/redoc'
    _PROXY_REDOC_URL = '/proxy-redoc'
    _OAUTH2_REDIRECT_URL = '/docs/oauth2-redirect'
    _PROXY_OAUTH2_REDIRECT_URL = '/proxy-docs/oauth2-redirect'

    # 文档静态资源URLs
    DEFAULT_REDOC_JS_URL = 'https://registry.npmmirror.com/redoc/2/files/bundles/redoc.standalone.js'
    DEFAULT_REDOC_FAVICON_URL = 'https://fastapi.tiangolo.com/img/favicon.png'
    DEFAULT_SWAGGER_JS_URL = 'https://registry.npmmirror.com/swagger-ui-dist/5/files/swagger-ui-bundle.js'
    DEFAULT_SWAGGER_CSS_URL = 'https://registry.npmmirror.com/swagger-ui-dist/5/files/swagger-ui.css'
    DEFAULT_SWAGGER_FAVICON_URL = 'https://fastapi.tiangolo.com/img/favicon.png'

    @classmethod
    def proxy_openapi_url(cls) -> str:
        """
        代理OpenAPI文档URL
        """
        return cls._PROXY_OPENAPI_URL if not AppConfig.app_disable_swagger and not AppConfig.app_disable_redoc else None

    @classmethod
    def docs_url(cls) -> str:
        """
        文档URL
        """
        return cls._DOCS_URL

    @classmethod
    def proxy_docs_url(cls) -> str:
        """
        代理文档URL
        """
        return cls._PROXY_DOCS_URL if not AppConfig.app_disable_swagger else None

    @classmethod
    def redoc_url(cls) -> str:
        """
        ReDoc文档URL
        """
        return cls._REDOC_URL

    @classmethod
    def proxy_redoc_url(cls) -> str:
        """
        代理ReDoc文档URL
        """
        return cls._PROXY_REDOC_URL if not AppConfig.app_disable_redoc else None

    @classmethod
    def proxy_oauth2_redirect_url(cls) -> str:
        """
        代理OAuth2重定向URL
        """
        return cls._PROXY_OAUTH2_REDIRECT_URL if not AppConfig.app_disable_swagger else None

    @classmethod
    def setup_docs_static_resources(
        cls,
        redoc_js_url: str = DEFAULT_REDOC_JS_URL,
        redoc_favicon_url: str = DEFAULT_REDOC_FAVICON_URL,
        swagger_js_url: str = DEFAULT_SWAGGER_JS_URL,
        swagger_css_url: str = DEFAULT_SWAGGER_CSS_URL,
        swagger_favicon_url: str = DEFAULT_SWAGGER_FAVICON_URL,
    ) -> None:
        """
        配置文档静态资源

        :param redoc_js_url: 用于加载ReDoc JavaScript的URL
        :param redoc_favicon_url: ReDoc要使用的favicon的URL
        :param swagger_js_url: 用于加载Swagger UI JavaScript的URL
        :param swagger_css_url: 用于加载Swagger UI CSS的URL
        :param swagger_favicon_url: Swagger UI要使用的favicon的URL
        :return:
        """

        def redoc_monkey_patch(*args, **kwargs) -> HTMLResponse:
            return get_redoc_html(
                *args,
                **kwargs,
                redoc_js_url=redoc_js_url,
                redoc_favicon_url=redoc_favicon_url,
            )

        def swagger_ui_monkey_patch(*args, **kwargs) -> HTMLResponse:
            return get_swagger_ui_html(
                *args,
                **kwargs,
                swagger_js_url=swagger_js_url,
                swagger_css_url=swagger_css_url,
                swagger_favicon_url=swagger_favicon_url,
            )

        applications.get_redoc_html = redoc_monkey_patch
        applications.get_swagger_ui_html = swagger_ui_monkey_patch

    @classmethod
    def custom_api_docs_router(
        cls,
        app: FastAPI,
        redoc_js_url: str = DEFAULT_REDOC_JS_URL,
        redoc_favicon_url: str = DEFAULT_REDOC_FAVICON_URL,
        swagger_js_url: str = DEFAULT_SWAGGER_JS_URL,
        swagger_css_url: str = DEFAULT_SWAGGER_CSS_URL,
        swagger_favicon_url: str = DEFAULT_SWAGGER_FAVICON_URL,
    ) -> None:
        """
        自定义API文档路由

        :param app: FastAPI对象
        :param redoc_js_url: 用于加载ReDoc JavaScript的URL
        :param redoc_favicon_url: ReDoc要使用的favicon的URL
        :param swagger_js_url: 用于加载Swagger UI JavaScript的URL
        :param swagger_css_url: 用于加载Swagger UI CSS的URL
        :param swagger_favicon_url: Swagger UI要使用的favicon的URL
        :return:
        """

        async def custom_openapi(request: Request) -> JSONResponse:
            return await cls._custom_openapi(app)

        async def custom_redoc(request: Request) -> HTMLResponse:
            return await cls._custom_redoc(app, redoc_js_url, redoc_favicon_url)

        async def custom_swagger(request: Request) -> HTMLResponse:
            return await cls._custom_swagger(app, swagger_js_url, swagger_css_url, swagger_favicon_url)

        async def custom_swagger_ui_redirect(request: Request) -> HTMLResponse:
            return await cls._custom_swagger_ui_redirect(app, swagger_favicon_url)

        # 注册路由
        app.add_route(cls._OPENAPI_URL, custom_openapi, include_in_schema=False)
        cls._register_docs_routes(app, custom_swagger, custom_swagger_ui_redirect, custom_redoc)

    @classmethod
    async def _custom_openapi(cls, app: FastAPI) -> JSONResponse:
        """
        自定义 OpenAPI 路由处理函数

        :param app: FastAPI对象
        :return: openapi的json响应
        """
        openapi_schema = get_openapi(
            title=app.title,
            version=app.version,
            openapi_version=app.openapi_version,
            summary=app.summary,
            description=app.description,
            terms_of_service=app.terms_of_service,
            contact=app.contact,
            license_info=app.license_info,
            routes=app.routes,
            webhooks=app.webhooks.routes,
            tags=app.openapi_tags,
            separate_input_output_schemas=app.separate_input_output_schemas,
            external_docs=app.openapi_external_docs,
        )
        return JSONResponse(openapi_schema)

    @classmethod
    async def _custom_redoc(
        cls,
        app: FastAPI,
        redoc_js_url: str,
        redoc_favicon_url: str,
    ) -> HTMLResponse:
        """
        自定义 ReDoc 路由处理函数

        :param app: FastAPI对象
        :param redoc_js_url: 用于加载ReDoc JavaScript的URL
        :param redoc_favicon_url: ReDoc要使用的favicon的URL
        :return: ReDoc HTML响应
        """
        if not AppConfig.app_disable_redoc:
            return get_redoc_html(
                openapi_url=cls._OPENAPI_URL,
                title=f'{app.title} - ReDoc',
                redoc_js_url=redoc_js_url,
                redoc_favicon_url=redoc_favicon_url,
            )
        return cls._get_disabled_html_content(
            f'{app.title} - ReDoc',
            'ReDoc',
            redoc_favicon_url,
        )

    @classmethod
    async def _custom_swagger(
        cls,
        app: FastAPI,
        swagger_js_url: str,
        swagger_css_url: str,
        swagger_favicon_url: str,
    ) -> HTMLResponse:
        """
        自定义 Swagger UI 路由处理函数

        :param app: FastAPI对象
        :param swagger_js_url: 用于加载Swagger UI JavaScript的URL
        :param swagger_css_url: 用于加载Swagger UI CSS的URL
        :param swagger_favicon_url: Swagger UI要使用的favicon的URL
        :return: Swagger UI HTML响应
        """
        if not AppConfig.app_disable_swagger:
            return get_swagger_ui_html(
                openapi_url=cls._OPENAPI_URL,
                title=f'{app.title} - Swagger UI',
                swagger_js_url=swagger_js_url,
                swagger_css_url=swagger_css_url,
                swagger_favicon_url=swagger_favicon_url,
                oauth2_redirect_url=app.swagger_ui_oauth2_redirect_url,
                init_oauth=app.swagger_ui_init_oauth,
                swagger_ui_parameters=app.swagger_ui_parameters,
            )
        return cls._get_disabled_html_content(
            f'{app.title} - Swagger UI',
            'Swagger UI',
            swagger_favicon_url,
        )

    @classmethod
    async def _custom_swagger_ui_redirect(
        cls,
        app: FastAPI,
        swagger_favicon_url: str,
    ) -> HTMLResponse:
        """
        自定义 Swagger UI OAuth2 重定向路由处理函数

        :param app: FastAPI对象
        :param swagger_favicon_url: Swagger UI要使用的favicon的URL
        :return: Swagger UI OAuth2重定向HTML响应
        """
        if not AppConfig.app_disable_swagger:
            return get_swagger_ui_oauth2_redirect_html()
        return cls._get_disabled_html_content(
            f'{app.title} - Swagger UI OAuth2 Redirect',
            'Swagger UI OAuth2 Redirect',
            swagger_favicon_url,
        )

    @staticmethod
    def _get_disabled_html_content(title: str, name: str, favicon_url: str) -> HTMLResponse:
        """
        生成禁用文档的HTML内容

        :param title: 页面标题
        :param name: 文档名称
        :param favicon_url: 图标地址
        :return: 禁用文档HTML响应
        """
        html = f"""
        <!DOCTYPE html>
        <html>
        <head>
        <title>{title}</title>
        <!-- needed for adaptive design -->
        <meta charset="utf-8"/>
        <meta name="viewport" content="width=device-width, initial-scale=1">
        <link rel="shortcut icon" href="{favicon_url}">
        <!--
        {name} doesn't change outer page styles
        -->
        <style>
        body {{
            margin: 10px;
            padding: 0;
        }}
        </style>
        </head>
        <body>
        <noscript>
            {name} requires Javascript to function. Please enable it to browse the documentation.
        </noscript>
        <h1 style="color: #ff4d4f;">{name} has been disabled. Please enable it first.</h1>
        </body>
        </html>
        """
        return HTMLResponse(html)

    @classmethod
    def _register_docs_routes(
        cls, app: FastAPI, swagger_handler: Callable, redirect_handler: Callable, redoc_handler: Callable
    ) -> None:
        """
        注册文档路由

        :param app: FastAPI对象
        :param swagger_handler: Swagger UI 路由处理函数
        :param redirect_handler: Swagger UI OAuth2 重定向路由处理函数
        :param redoc_handler: ReDoc 路由处理函数
        :return:
        """
        swagger_urls: list[str] = (
            [cls._DOCS_URL] if not AppConfig.app_disable_swagger else [cls._DOCS_URL, cls._PROXY_DOCS_URL]
        )
        swagger_redirect_urls: list[str] = (
            [cls._OAUTH2_REDIRECT_URL]
            if not AppConfig.app_disable_swagger
            else [cls._OAUTH2_REDIRECT_URL, cls._PROXY_OAUTH2_REDIRECT_URL]
        )
        redoc_urls: list[str] = (
            [cls._REDOC_URL] if not AppConfig.app_disable_redoc else [cls._REDOC_URL, cls._PROXY_REDOC_URL]
        )

        for url in swagger_urls:
            app.add_route(url, swagger_handler, include_in_schema=False)
        for url in swagger_redirect_urls:
            app.add_route(url, redirect_handler, include_in_schema=False)
        for url in redoc_urls:
            app.add_route(url, redoc_handler, include_in_schema=False)


class StartupUtil:
    """
    启动门禁工具类
    """

    @classmethod
    async def acquire_startup_log_gate(
        cls, redis: aioredis.Redis, lock_key: str, worker_id: str, lock_expire_seconds: int
    ) -> bool:
        """
        获取启动日志门禁

        :param redis: Redis连接对象
        :param lock_key: 分布式锁key
        :param worker_id: 当前worker标识
        :param lock_expire_seconds: 锁过期时间
        :return: 是否获得启动日志输出权
        """
        acquired = await redis.set(lock_key, worker_id, nx=True, ex=lock_expire_seconds)
        if acquired:
            return True
        current_holder = await redis.get(lock_key)
        return current_holder == worker_id

    @classmethod
    def start_lock_renewal(
        cls,
        redis: aioredis.Redis,
        lock_key: str,
        worker_id: str,
        lock_expire_seconds: int,
        interval_seconds: int,
        on_lock_lost: Callable[[], None] | None = None,
    ) -> asyncio.Task:
        """
        启动分布式锁续期任务

        :param redis: Redis连接对象
        :param lock_key: 分布式锁key
        :param worker_id: 当前worker标识
        :param lock_expire_seconds: 锁过期时间
        :param interval_seconds: 续期间隔时间
        :param on_lock_lost: 失去锁时的回调
        :return: 异步任务对象
        """

        async def _loop() -> None:
            while True:
                try:
                    current_holder = await redis.get(lock_key)
                    if current_holder == worker_id:
                        await redis.expire(lock_key, lock_expire_seconds)
                        await asyncio.sleep(interval_seconds)
                        continue
                    if on_lock_lost:
                        on_lock_lost()
                    break
                except Exception:
                    await asyncio.sleep(interval_seconds)

        return asyncio.create_task(_loop())


class WorkerIdUtil:
    """
    Worker标识生成工具类
    """

    _worker_id: str | None = None

    @classmethod
    def get_worker_id(cls, configured_worker_id: str | None) -> str:
        """
        获取当前worker标识

        :param configured_worker_id: 配置的worker标识
        :return: 当前worker标识
        """
        if cls._worker_id:
            return cls._worker_id
        worker_id = configured_worker_id
        if not worker_id or worker_id.lower() == 'auto':
            worker_id = f'{os.getpid()}-{uuid.uuid4().hex[:6]}'
        cls._worker_id = worker_id
        return worker_id


class IPUtil:
    """
    IP工具类
    """

    _PREFERRED_DNS_HOSTS: tuple[str, str] = ('223.5.5.5', '8.8.8.8')
    _DNS_CONNECT_TIMEOUT = 1

    @classmethod
    def get_local_ip(cls) -> str:
        """
        获取本机Local IP
        """
        try:
            for snics in psutil.net_if_addrs().values():
                for snic in snics:
                    if snic.family == socket.AF_INET and snic.address.startswith('127.'):
                        return snic.address
        except Exception:
            pass

        return '127.0.0.1'

    @classmethod
    def get_network_ips(cls) -> list[str]:
        """
        获取本机Network IP列表
        """
        network_ips = []
        try:
            # 获取网卡状态
            stats = psutil.net_if_stats()
            for name, snics in psutil.net_if_addrs().items():
                # 过滤掉状态为DOWN的网卡
                if name in stats and not stats[name].isup:
                    continue

                for snic in snics:
                    if snic.family == socket.AF_INET:
                        try:
                            ip_obj = ipaddress.ip_address(snic.address)
                            if ip_obj.is_loopback or ip_obj.is_link_local:
                                continue
                            network_ips.append(snic.address)
                        except ValueError:
                            continue
        except Exception:
            pass

        # 优先显示首选出站IP
        preferred_ip = None
        for dns_host in cls._PREFERRED_DNS_HOSTS:
            try:
                with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
                    s.settimeout(cls._DNS_CONNECT_TIMEOUT)
                    s.connect((dns_host, 80))
                    preferred_ip = s.getsockname()[0]
                    break
            except Exception:
                continue

        if preferred_ip:
            if preferred_ip in network_ips:
                network_ips.remove(preferred_ip)
            network_ips.insert(0, preferred_ip)

        if not network_ips:
            network_ips = ['127.0.0.1']

        return network_ips