176 lines
6.2 KiB
Python
176 lines
6.2 KiB
Python
"""多 IDP 信任边界注册表 (R1a, U4 SSO 集成)。
|
||
|
||
管理多个 IDP 配置（OIDC + SAML 并存，如飞书 / Azure AD / 阿里云 IDaaS）：
|
||
|
||
- **热证书轮换**：``update_certificate`` 直接替换内存配置，无需重启。
|
||
- **单 IDP 故障隔离**：某个 IDP 不可用不影响其他 IDP 的认证流程。
|
||
- **按 sub/NameID 关联**（R1a）：禁止仅凭邮箱合并账户 — 见 :mod:`providers.oidc`。
|
||
|
||
配置来源：``agentkit.yaml`` 的 ``auth.idps`` 段（见 :func:`IDPRegistry.from_config`）。
|
||
|
||
单租户假设 (R1c)：所有 IDP 共享同一本地用户表，无租户隔离。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import threading
|
||
from typing import Literal
|
||
|
||
from pydantic import BaseModel, ConfigDict, Field
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
IDPType = Literal["oidc", "saml"]
|
||
|
||
|
||
class OIDCConfig(BaseModel):
|
||
"""单个 OIDC IDP 配置（authlib redirect flow）。"""
|
||
|
||
model_config = ConfigDict(extra="forbid")
|
||
|
||
client_id: str
|
||
client_secret: str
|
||
server_metadata_url: str
|
||
# redirect_uri 是完整 URL，如 https://host/api/v1/auth/oauth/feishu/callback
|
||
redirect_uri: str
|
||
# 可选：scope 覆盖（默认 openid email profile）
|
||
scope: str = "openid email profile"
|
||
|
||
|
||
class SAMLConfig(BaseModel):
|
||
"""单个 SAML 2.0 IDP 配置（python3-saml / OneLogin）。"""
|
||
|
||
model_config = ConfigDict(extra="forbid")
|
||
|
||
# IDP 元数据：entity_id + 单点登录 URL（SP 发起时构造 AuthnRequest）
|
||
entity_id: str
|
||
sso_url: str
|
||
# IDP 签名证书（PEM 格式，可能多张证书用于轮换）
|
||
idp_cert: str
|
||
# SP 自身配置
|
||
sp_entity_id: str
|
||
acs_url: str
|
||
# 可选：SLO URL
|
||
slo_url: str | None = None
|
||
|
||
|
||
class IDPConfig(BaseModel):
|
||
"""单个 IDP 完整配置（OIDC 或 SAML 之一）。"""
|
||
|
||
model_config = ConfigDict(extra="forbid")
|
||
|
||
name: str = Field(..., description="IDP 唯一标识，如 feishu / azure_ad / aliyun_idaas")
|
||
type: IDPType
|
||
display_name: str = ""
|
||
enabled: bool = True
|
||
oidc: OIDCConfig | None = None
|
||
saml: SAMLConfig | None = None
|
||
|
||
def is_healthy(self) -> bool:
|
||
"""IDP 是否可用（启用 + 配置完整）。
|
||
|
||
单 IDP 故障不影响其他 — 调用方逐个检查此方法后决定是否跳过。
|
||
"""
|
||
if not self.enabled:
|
||
return False
|
||
if self.type == "oidc":
|
||
return self.oidc is not None
|
||
return self.saml is not None
|
||
|
||
|
||
class IDPRegistry:
|
||
"""多 IDP 注册表 — 进程内单例，支持热证书轮换。
|
||
|
||
线程安全：所有写操作加锁（``_lock``）。读操作返回配置的浅拷贝引用 —
|
||
配置对象本身不可变（Pydantic），热轮换通过整体替换实现。
|
||
|
||
ponytail: 进程内 dict 注册表，不持久化。多实例部署时各进程独立加载
|
||
agentkit.yaml；证书热轮换需配合配置同步（config_sync.py 轮询）或重启。
|
||
上限：单进程数百 IDP 无压力（O(1) 查找）。
|
||
"""
|
||
|
||
def __init__(self) -> None:
|
||
self._idps: dict[str, IDPConfig] = {}
|
||
self._lock = threading.Lock()
|
||
|
||
def register(self, config: IDPConfig) -> None:
|
||
"""注册或覆盖一个 IDP（用于热加载 / 热证书轮换）。"""
|
||
with self._lock:
|
||
self._idps[config.name] = config
|
||
logger.info(f"IDP 已注册: {config.name} (type={config.type}, enabled={config.enabled})")
|
||
|
||
def remove(self, name: str) -> bool:
|
||
"""移除一个 IDP。返回是否实际移除。"""
|
||
with self._lock:
|
||
return self._idps.pop(name, None) is not None
|
||
|
||
def get(self, name: str) -> IDPConfig | None:
|
||
"""按名称查找 IDP 配置。返回 ``None`` 表示不存在。"""
|
||
with self._lock:
|
||
return self._idps.get(name)
|
||
|
||
def list_idps(self) -> list[IDPConfig]:
|
||
"""列出所有已注册 IDP（含禁用的）。"""
|
||
with self._lock:
|
||
return list(self._idps.values())
|
||
|
||
def list_healthy(self) -> list[IDPConfig]:
|
||
"""列出所有可用 IDP（启用 + 配置完整）。"""
|
||
return [c for c in self.list_idps() if c.is_healthy()]
|
||
|
||
def update_certificate(self, name: str, *, idp_cert: str | None = None) -> bool:
|
||
"""热轮换 SAML IDP 签名证书 — 不重启即时生效。
|
||
|
||
Args:
|
||
name: IDP 名称。
|
||
idp_cert: 新的 IDP 签名证书（PEM）。
|
||
|
||
Returns:
|
||
``True`` 表示成功更新，``False`` 表示 IDP 不存在或非 SAML 类型。
|
||
|
||
OIDC 证书轮换通过 ``server_metadata_url`` 自动发现（authlib 缓存 TTL），
|
||
无需显式热轮换 — 调用 ``register`` 替换整个 :class:`OIDCConfig` 即可。
|
||
"""
|
||
with self._lock:
|
||
existing = self._idps.get(name)
|
||
if existing is None or existing.type != "saml" or existing.saml is None:
|
||
return False
|
||
# Pydantic 不可变 — 构造新配置整体替换
|
||
new_saml = existing.saml.model_copy(update={"idp_cert": idp_cert or ""})
|
||
self._idps[name] = existing.model_copy(update={"saml": new_saml})
|
||
logger.info(f"IDP {name} 签名证书已热轮换")
|
||
return True
|
||
|
||
def load_from_config(self, idps: list[dict[str, object]]) -> None:
|
||
"""从 agentkit.yaml 的 ``auth.idps`` 列表批量加载（覆盖现有）。"""
|
||
with self._lock:
|
||
self._idps.clear()
|
||
for raw in idps:
|
||
try:
|
||
cfg = IDPConfig.model_validate(raw)
|
||
self.register(cfg)
|
||
except Exception as exc: # noqa: BLE001 — 单个 IDP 配置错误不阻断其他
|
||
logger.error(f"IDP 配置加载失败，跳过: {raw.get('name', '?')}: {exc}")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 进程级单例
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_registry: IDPRegistry | None = None
|
||
|
||
|
||
def get_idp_registry() -> IDPRegistry:
|
||
"""返回进程级 IDP 注册表单例（惰性初始化）。"""
|
||
global _registry
|
||
if _registry is None:
|
||
_registry = IDPRegistry()
|
||
return _registry
|
||
|
||
|
||
def reset_idp_registry() -> None:
|
||
"""重置注册表单例（测试用）。"""
|
||
global _registry
|
||
_registry = None
|