24024
af7c11d7f9
- 新增图像生成接口，支持试用、积分和自定义API Key模式 - 实现生成图片结果异步上传至MinIO存储，带重试机制 - 优化积分预扣除和异常退还逻辑，保障用户积分准确 - 添加获取生成历史记录接口，支持时间范围和分页 - 提供本地字典配置接口，支持模型、比例、提示模板和尺寸 - 实现图片批量上传接口，支持S3兼容对象存储 feat(admin): 增加管理员角色管理与权限分配接口 - 实现角色列表查询、角色创建、更新及删除功能 - 增加权限列表查询接口 - 实现用户角色分配接口，便于统一管理用户权限 - 增加系统字典增删查改接口，支持分类过滤和排序 - 权限控制全面覆盖管理接口，保证安全访问 feat(auth): 完善用户登录注册及权限相关接口与页面 - 实现手机号验证码发送及校验功能，保障注册安全 - 支持手机号注册、登录及退出接口，集成日志记录 - 增加修改密码功能，验证原密码后更新 - 提供动态导航菜单接口，基于权限展示不同菜单 - 实现管理界面路由及日志、角色、字典管理页面访问权限控制 - 添加系统日志查询接口，支持关键词和等级筛选 feat(app): 初始化Flask应用并配置蓝图与数据库 - 创建应用程序工厂，加载配置，初始化数据库和Redis客户端 - 注册认证、API及管理员蓝图，整合路由 - 根路由渲染主页模板 - 应用上下文中自动创建数据库表，保证运行环境准备完毕 feat(database): 提供数据库创建与迁移支持脚本 - 新增数据库创建脚本，支持自动检测是否已存在 - 添加数据库表初始化脚本，支持创建和删除所有表 - 实现RBAC权限初始化，包含基础权限和角色创建 - 新增字段手动修复脚本，添加用户API Key和积分字段 - 强制迁移脚本支持清理连接和修复表结构，初始化默认数据及角色分配 feat(config): 新增系统配置参数 - 配置数据库、Redis、Session和MinIO相关参数 - 添加AI接口地址及试用Key配置 - 集成阿里云短信服务配置及开发模式相关参数 feat(extensions): 初始化数据库、Redis和MinIO客户端 - 创建全局SQLAlchemy数据库实例和Redis客户端 - 配置基于boto3的MinIO兼容S3客户端 chore(logs): 添加示例系统日志文件 - 记录用户请求、验证码发送成功与失败的日志信息
371 lines
12 KiB
Python
371 lines
12 KiB
Python
import asyncio
|
|
import logging
|
|
import threading
|
|
from datetime import datetime, timezone
|
|
from time import sleep
|
|
from typing import Any, Awaitable, Callable, Union
|
|
|
|
from redis.auth.err import RequestTokenErr, TokenRenewalErr
|
|
from redis.auth.idp import IdentityProviderInterface
|
|
from redis.auth.token import TokenResponse
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class CredentialsListener:
|
|
"""
|
|
Listeners that will be notified on events related to credentials.
|
|
Accepts callbacks and awaitable callbacks.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._on_next = None
|
|
self._on_error = None
|
|
|
|
@property
|
|
def on_next(self) -> Union[Callable[[Any], None], Awaitable]:
|
|
return self._on_next
|
|
|
|
@on_next.setter
|
|
def on_next(self, callback: Union[Callable[[Any], None], Awaitable]) -> None:
|
|
self._on_next = callback
|
|
|
|
@property
|
|
def on_error(self) -> Union[Callable[[Exception], None], Awaitable]:
|
|
return self._on_error
|
|
|
|
@on_error.setter
|
|
def on_error(self, callback: Union[Callable[[Exception], None], Awaitable]) -> None:
|
|
self._on_error = callback
|
|
|
|
|
|
class RetryPolicy:
|
|
def __init__(self, max_attempts: int, delay_in_ms: float):
|
|
self.max_attempts = max_attempts
|
|
self.delay_in_ms = delay_in_ms
|
|
|
|
def get_max_attempts(self) -> int:
|
|
"""
|
|
Retry attempts before exception will be thrown.
|
|
|
|
:return: int
|
|
"""
|
|
return self.max_attempts
|
|
|
|
def get_delay_in_ms(self) -> float:
|
|
"""
|
|
Delay between retries in seconds.
|
|
|
|
:return: int
|
|
"""
|
|
return self.delay_in_ms
|
|
|
|
|
|
class TokenManagerConfig:
|
|
def __init__(
|
|
self,
|
|
expiration_refresh_ratio: float,
|
|
lower_refresh_bound_millis: int,
|
|
token_request_execution_timeout_in_ms: int,
|
|
retry_policy: RetryPolicy,
|
|
):
|
|
self._expiration_refresh_ratio = expiration_refresh_ratio
|
|
self._lower_refresh_bound_millis = lower_refresh_bound_millis
|
|
self._token_request_execution_timeout_in_ms = (
|
|
token_request_execution_timeout_in_ms
|
|
)
|
|
self._retry_policy = retry_policy
|
|
|
|
def get_expiration_refresh_ratio(self) -> float:
|
|
"""
|
|
Represents the ratio of a token's lifetime at which a refresh should be triggered. # noqa: E501
|
|
For example, a value of 0.75 means the token should be refreshed
|
|
when 75% of its lifetime has elapsed (or when 25% of its lifetime remains).
|
|
|
|
:return: float
|
|
"""
|
|
|
|
return self._expiration_refresh_ratio
|
|
|
|
def get_lower_refresh_bound_millis(self) -> int:
|
|
"""
|
|
Represents the minimum time in milliseconds before token expiration
|
|
to trigger a refresh, in milliseconds.
|
|
This value sets a fixed lower bound for when a token refresh should occur,
|
|
regardless of the token's total lifetime.
|
|
If set to 0 there will be no lower bound and the refresh will be triggered
|
|
based on the expirationRefreshRatio only.
|
|
|
|
:return: int
|
|
"""
|
|
return self._lower_refresh_bound_millis
|
|
|
|
def get_token_request_execution_timeout_in_ms(self) -> int:
|
|
"""
|
|
Represents the maximum time in milliseconds to wait
|
|
for a token request to complete.
|
|
|
|
:return: int
|
|
"""
|
|
return self._token_request_execution_timeout_in_ms
|
|
|
|
def get_retry_policy(self) -> RetryPolicy:
|
|
"""
|
|
Represents the retry policy for token requests.
|
|
|
|
:return: RetryPolicy
|
|
"""
|
|
return self._retry_policy
|
|
|
|
|
|
class TokenManager:
|
|
def __init__(
|
|
self, identity_provider: IdentityProviderInterface, config: TokenManagerConfig
|
|
):
|
|
self._idp = identity_provider
|
|
self._config = config
|
|
self._next_timer = None
|
|
self._listener = None
|
|
self._init_timer = None
|
|
self._retries = 0
|
|
|
|
def __del__(self):
|
|
logger.info("Token manager are disposed")
|
|
self.stop()
|
|
|
|
def start(
|
|
self,
|
|
listener: CredentialsListener,
|
|
skip_initial: bool = False,
|
|
) -> Callable[[], None]:
|
|
self._listener = listener
|
|
|
|
try:
|
|
loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
# Run loop in a separate thread to unblock main thread.
|
|
loop = asyncio.new_event_loop()
|
|
thread = threading.Thread(
|
|
target=_start_event_loop_in_thread, args=(loop,), daemon=True
|
|
)
|
|
thread.start()
|
|
|
|
# Event to block for initial execution.
|
|
init_event = asyncio.Event()
|
|
self._init_timer = loop.call_later(
|
|
0, self._renew_token, skip_initial, init_event
|
|
)
|
|
logger.info("Token manager started")
|
|
|
|
# Blocks in thread-safe manner.
|
|
asyncio.run_coroutine_threadsafe(init_event.wait(), loop).result()
|
|
return self.stop
|
|
|
|
async def start_async(
|
|
self,
|
|
listener: CredentialsListener,
|
|
block_for_initial: bool = False,
|
|
initial_delay_in_ms: float = 0,
|
|
skip_initial: bool = False,
|
|
) -> Callable[[], None]:
|
|
self._listener = listener
|
|
|
|
loop = asyncio.get_running_loop()
|
|
init_event = asyncio.Event()
|
|
|
|
# Wraps the async callback with async wrapper to schedule with loop.call_later()
|
|
wrapped = _async_to_sync_wrapper(
|
|
loop, self._renew_token_async, skip_initial, init_event
|
|
)
|
|
self._init_timer = loop.call_later(initial_delay_in_ms / 1000, wrapped)
|
|
logger.info("Token manager started")
|
|
|
|
if block_for_initial:
|
|
await init_event.wait()
|
|
|
|
return self.stop
|
|
|
|
def stop(self):
|
|
if self._init_timer is not None:
|
|
self._init_timer.cancel()
|
|
if self._next_timer is not None:
|
|
self._next_timer.cancel()
|
|
|
|
def acquire_token(self, force_refresh=False) -> TokenResponse:
|
|
try:
|
|
token = self._idp.request_token(force_refresh)
|
|
except RequestTokenErr as e:
|
|
if self._retries < self._config.get_retry_policy().get_max_attempts():
|
|
self._retries += 1
|
|
sleep(self._config.get_retry_policy().get_delay_in_ms() / 1000)
|
|
return self.acquire_token(force_refresh)
|
|
else:
|
|
raise e
|
|
|
|
self._retries = 0
|
|
return TokenResponse(token)
|
|
|
|
async def acquire_token_async(self, force_refresh=False) -> TokenResponse:
|
|
try:
|
|
token = self._idp.request_token(force_refresh)
|
|
except RequestTokenErr as e:
|
|
if self._retries < self._config.get_retry_policy().get_max_attempts():
|
|
self._retries += 1
|
|
await asyncio.sleep(
|
|
self._config.get_retry_policy().get_delay_in_ms() / 1000
|
|
)
|
|
return await self.acquire_token_async(force_refresh)
|
|
else:
|
|
raise e
|
|
|
|
self._retries = 0
|
|
return TokenResponse(token)
|
|
|
|
def _calculate_renewal_delay(self, expire_date: float, issue_date: float) -> float:
|
|
delay_for_lower_refresh = self._delay_for_lower_refresh(expire_date)
|
|
delay_for_ratio_refresh = self._delay_for_ratio_refresh(expire_date, issue_date)
|
|
delay = min(delay_for_ratio_refresh, delay_for_lower_refresh)
|
|
|
|
return 0 if delay < 0 else delay / 1000
|
|
|
|
def _delay_for_lower_refresh(self, expire_date: float):
|
|
return (
|
|
expire_date
|
|
- self._config.get_lower_refresh_bound_millis()
|
|
- (datetime.now(timezone.utc).timestamp() * 1000)
|
|
)
|
|
|
|
def _delay_for_ratio_refresh(self, expire_date: float, issue_date: float):
|
|
token_ttl = expire_date - issue_date
|
|
refresh_before = token_ttl - (
|
|
token_ttl * self._config.get_expiration_refresh_ratio()
|
|
)
|
|
|
|
return (
|
|
expire_date
|
|
- refresh_before
|
|
- (datetime.now(timezone.utc).timestamp() * 1000)
|
|
)
|
|
|
|
def _renew_token(
|
|
self, skip_initial: bool = False, init_event: asyncio.Event = None
|
|
):
|
|
"""
|
|
Task to renew token from identity provider.
|
|
Schedules renewal tasks based on token TTL.
|
|
"""
|
|
|
|
try:
|
|
token_res = self.acquire_token(force_refresh=True)
|
|
delay = self._calculate_renewal_delay(
|
|
token_res.get_token().get_expires_at_ms(),
|
|
token_res.get_token().get_received_at_ms(),
|
|
)
|
|
|
|
if token_res.get_token().is_expired():
|
|
raise TokenRenewalErr("Requested token is expired")
|
|
|
|
if self._listener.on_next is None:
|
|
logger.warning(
|
|
"No registered callback for token renewal task. Renewal cancelled"
|
|
)
|
|
return
|
|
|
|
if not skip_initial:
|
|
try:
|
|
self._listener.on_next(token_res.get_token())
|
|
except Exception as e:
|
|
raise TokenRenewalErr(e)
|
|
|
|
if delay <= 0:
|
|
return
|
|
|
|
loop = asyncio.get_running_loop()
|
|
self._next_timer = loop.call_later(delay, self._renew_token)
|
|
logger.info(f"Next token renewal scheduled in {delay} seconds")
|
|
return token_res
|
|
except Exception as e:
|
|
if self._listener.on_error is None:
|
|
raise e
|
|
|
|
self._listener.on_error(e)
|
|
finally:
|
|
if init_event:
|
|
init_event.set()
|
|
|
|
async def _renew_token_async(
|
|
self, skip_initial: bool = False, init_event: asyncio.Event = None
|
|
):
|
|
"""
|
|
Async task to renew tokens from identity provider.
|
|
Schedules renewal tasks based on token TTL.
|
|
"""
|
|
|
|
try:
|
|
token_res = await self.acquire_token_async(force_refresh=True)
|
|
delay = self._calculate_renewal_delay(
|
|
token_res.get_token().get_expires_at_ms(),
|
|
token_res.get_token().get_received_at_ms(),
|
|
)
|
|
|
|
if token_res.get_token().is_expired():
|
|
raise TokenRenewalErr("Requested token is expired")
|
|
|
|
if self._listener.on_next is None:
|
|
logger.warning(
|
|
"No registered callback for token renewal task. Renewal cancelled"
|
|
)
|
|
return
|
|
|
|
if not skip_initial:
|
|
try:
|
|
await self._listener.on_next(token_res.get_token())
|
|
except Exception as e:
|
|
raise TokenRenewalErr(e)
|
|
|
|
if delay <= 0:
|
|
return
|
|
|
|
loop = asyncio.get_running_loop()
|
|
wrapped = _async_to_sync_wrapper(loop, self._renew_token_async)
|
|
logger.info(f"Next token renewal scheduled in {delay} seconds")
|
|
loop.call_later(delay, wrapped)
|
|
except Exception as e:
|
|
if self._listener.on_error is None:
|
|
raise e
|
|
|
|
await self._listener.on_error(e)
|
|
finally:
|
|
if init_event:
|
|
init_event.set()
|
|
|
|
|
|
def _async_to_sync_wrapper(loop, coro_func, *args, **kwargs):
|
|
"""
|
|
Wraps an asynchronous function so it can be used with loop.call_later.
|
|
|
|
:param loop: The event loop in which the coroutine will be executed.
|
|
:param coro_func: The coroutine function to wrap.
|
|
:param args: Positional arguments to pass to the coroutine function.
|
|
:param kwargs: Keyword arguments to pass to the coroutine function.
|
|
:return: A regular function suitable for loop.call_later.
|
|
"""
|
|
|
|
def wrapped():
|
|
# Schedule the coroutine in the event loop
|
|
asyncio.ensure_future(coro_func(*args, **kwargs), loop=loop)
|
|
|
|
return wrapped
|
|
|
|
|
|
def _start_event_loop_in_thread(event_loop: asyncio.AbstractEventLoop):
|
|
"""
|
|
Starts event loop in a thread.
|
|
Used to be able to schedule tasks using loop.call_later.
|
|
|
|
:param event_loop:
|
|
:return:
|
|
"""
|
|
asyncio.set_event_loop(event_loop)
|
|
event_loop.run_forever()
|