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""" {title}

{name} has been disabled. Please enable it first.

""" 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