Kris
655eab7230
新增了完整的认证工具库模块，包含以下功能： 1. 指数退避重试组件 2. 敏感数据过滤与日志脱敏 3. 安全策略管理引擎 4. 认证健康监控模块 5. SSRF防护工具集 6. 多类型token管理系统 7. 密钥管理器加解密工具
141 lines
4.4 KiB
Python
141 lines
4.4 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
|
|
@dataclass
|
|
class AuthHealth:
|
|
status: str
|
|
message: str = ""
|
|
expires_at: float | None = None
|
|
last_checked_at: float = field(default_factory=time.monotonic)
|
|
metadata: dict[str, Any] = field(default_factory=dict)
|
|
|
|
@property
|
|
def is_healthy(self) -> bool:
|
|
return self.status == "healthy"
|
|
|
|
|
|
class AuthHealthMonitor:
|
|
def __init__(self):
|
|
self._health_states: dict[str, AuthHealth] = {}
|
|
self._check_tasks: dict[str, asyncio.Task] = {}
|
|
self._on_expiring_callbacks: list = []
|
|
self._on_expired_callbacks: list = []
|
|
self._on_revoked_callbacks: list = []
|
|
self._on_failure_callbacks: list = []
|
|
self._locks: dict[str, asyncio.Lock] = {}
|
|
|
|
def _get_lock(self, channel_id: str) -> asyncio.Lock:
|
|
if channel_id not in self._locks:
|
|
self._locks[channel_id] = asyncio.Lock()
|
|
return self._locks[channel_id]
|
|
|
|
def update_health(self, channel_id: str, health: AuthHealth) -> None:
|
|
old = self._health_states.get(channel_id)
|
|
self._health_states[channel_id] = health
|
|
|
|
if old and old.status == "healthy" and health.status != "healthy":
|
|
try:
|
|
asyncio.create_task(self._notify_failure(channel_id, health))
|
|
except RuntimeError:
|
|
pass
|
|
|
|
def get_health(self, channel_id: str) -> AuthHealth | None:
|
|
return self._health_states.get(channel_id)
|
|
|
|
def get_all_health(self) -> dict[str, AuthHealth]:
|
|
return dict(self._health_states)
|
|
|
|
def on_token_expiring(self, callback) -> None:
|
|
self._on_expiring_callbacks.append(callback)
|
|
|
|
def on_token_expired(self, callback) -> None:
|
|
self._on_expired_callbacks.append(callback)
|
|
|
|
def on_auth_revoked(self, callback) -> None:
|
|
self._on_revoked_callbacks.append(callback)
|
|
|
|
def on_auth_failure(self, callback) -> None:
|
|
self._on_failure_callbacks.append(callback)
|
|
|
|
async def _notify_expiring(self, channel_id: str) -> None:
|
|
for cb in self._on_expiring_callbacks:
|
|
try:
|
|
await cb(channel_id)
|
|
except Exception:
|
|
pass
|
|
|
|
async def _notify_expired(self, channel_id: str) -> None:
|
|
for cb in self._on_expired_callbacks:
|
|
try:
|
|
await cb(channel_id)
|
|
except Exception:
|
|
pass
|
|
|
|
async def _notify_revoked(self, channel_id: str) -> None:
|
|
for cb in self._on_revoked_callbacks:
|
|
try:
|
|
await cb(channel_id)
|
|
except Exception:
|
|
pass
|
|
|
|
async def _notify_failure(self, channel_id: str, health: AuthHealth) -> None:
|
|
for cb in self._on_failure_callbacks:
|
|
try:
|
|
await cb(channel_id, health)
|
|
except Exception:
|
|
pass
|
|
|
|
def start_periodic_check(
|
|
self,
|
|
channel_id: str,
|
|
check_func,
|
|
interval_seconds: float = 60.0,
|
|
) -> None:
|
|
if channel_id in self._check_tasks and not self._check_tasks[channel_id].done():
|
|
return
|
|
self._check_tasks[channel_id] = asyncio.create_task(
|
|
self._periodic_check_loop(channel_id, check_func, interval_seconds)
|
|
)
|
|
|
|
def stop_periodic_check(self, channel_id: str) -> None:
|
|
task = self._check_tasks.pop(channel_id, None)
|
|
if task and not task.done():
|
|
task.cancel()
|
|
|
|
async def _periodic_check_loop(
|
|
self,
|
|
channel_id: str,
|
|
check_func,
|
|
interval_seconds: float,
|
|
) -> None:
|
|
from yuxi.utils.logging_config import logger
|
|
|
|
while True:
|
|
try:
|
|
await asyncio.sleep(interval_seconds)
|
|
health = await check_func()
|
|
self.update_health(channel_id, health)
|
|
except asyncio.CancelledError:
|
|
break
|
|
except Exception as e:
|
|
logger.warning(f"[AuthHealthMonitor] Health check failed for {channel_id}: {e}")
|
|
self.update_health(
|
|
channel_id,
|
|
AuthHealth(status="unhealthy", message=str(e)),
|
|
)
|
|
|
|
|
|
_health_monitor: AuthHealthMonitor | None = None
|
|
|
|
|
|
def get_health_monitor() -> AuthHealthMonitor:
|
|
global _health_monitor
|
|
if _health_monitor is None:
|
|
_health_monitor = AuthHealthMonitor()
|
|
return _health_monitor
|