Kris
1f78c44b03
这是一个批量整理提交，包含以下主要改动： 1. 删除多处冗余的空行和未使用的导入 2. 修复文件末尾缺少换行符的问题 3. 调整部分模块的导入顺序与代码排版 4. 修复部分配置默认值与策略逻辑 5. 新增多个功能模块与辅助工具 6. 完善异常处理与日志记录 7. 修复速率限制、消息缓存、权限校验等逻辑bug 8. 废弃部分旧有API与配置项并添加警告提示
261 lines
9.5 KiB
Python
261 lines
9.5 KiB
Python
"""Microsoft Teams 凭据管理。
|
|
|
|
支持 Secret、Certificate (Federated) 和 Managed Identity 三种凭据模式，
|
|
以及 Delegated Auth (OAuth 2.0 Refresh Token)。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import aiohttp
|
|
|
|
from yuxi.utils.logging_config import logger
|
|
|
|
TOKEN_URL_TEMPLATE = "https://login.microsoftonline.com/{tenant}/oauth2/v2.0/token"
|
|
BOT_SCOPE = "https://api.botframework.com/.default"
|
|
GRAPH_SCOPE = "https://graph.microsoft.com/.default"
|
|
TOKEN_GRACE_PERIOD_S = 60
|
|
|
|
DELEGATED_STORE_FILENAME = "msteams-delegated.json"
|
|
DELEGATED_REFRESH_BUFFER_S = 300
|
|
|
|
|
|
class FederatedCredentialError(Exception):
|
|
"""联邦凭证错误，区分于其他认证错误。"""
|
|
|
|
pass
|
|
|
|
|
|
class FederatedCredential:
|
|
def __init__(
|
|
self,
|
|
client_id: str,
|
|
tenant_id: str = "common",
|
|
certificate_path: str = "",
|
|
certificate_thumbprint: str = "",
|
|
use_managed_identity: bool = False,
|
|
):
|
|
self.client_id = client_id
|
|
self.tenant_id = tenant_id
|
|
self.certificate_path = certificate_path
|
|
self.certificate_thumbprint = certificate_thumbprint
|
|
self.use_managed_identity = use_managed_identity
|
|
if self.certificate_path and not self.use_managed_identity:
|
|
self._check_prerequisites()
|
|
|
|
@property
|
|
def auth_type(self) -> str:
|
|
if self.use_managed_identity:
|
|
return "managed_identity"
|
|
if self.certificate_path:
|
|
return "certificate"
|
|
return "secret"
|
|
|
|
def _check_prerequisites(self) -> None:
|
|
try:
|
|
from cryptography.hazmat.primitives import serialization # noqa: F401
|
|
from cryptography.hazmat.backends import default_backend # noqa: F401
|
|
import jwt as pyjwt # noqa: F401
|
|
except ImportError as e:
|
|
raise FederatedCredentialError(
|
|
"FederatedCredential requires 'cryptography' and 'PyJWT'. Install with: pip install cryptography pyjwt"
|
|
) from e
|
|
|
|
cert_path = Path(self.certificate_path)
|
|
if not cert_path.exists():
|
|
raise FederatedCredentialError(f"Certificate file not found: {self.certificate_path}")
|
|
|
|
async def get_token(self, scope: str) -> str | None:
|
|
if self.use_managed_identity:
|
|
return await self._get_managed_identity_token(scope)
|
|
if self.certificate_path:
|
|
return await self._get_certificate_token(scope)
|
|
return None
|
|
|
|
async def _get_managed_identity_token(self, scope: str) -> str | None:
|
|
endpoint = os.getenv("MSI_ENDPOINT", "http://169.254.169.254/metadata/identity/oauth2/token")
|
|
params = {
|
|
"api-version": "2019-08-01",
|
|
"resource": scope.rstrip("/.default"),
|
|
"client_id": self.client_id,
|
|
}
|
|
try:
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.get(
|
|
endpoint,
|
|
params=params,
|
|
headers={"Metadata": "true"},
|
|
) as resp:
|
|
if resp.status == 200:
|
|
result = await resp.json()
|
|
return result.get("access_token")
|
|
logger.warning(f"Managed Identity token failed: HTTP {resp.status}")
|
|
except Exception as e:
|
|
logger.warning(f"Managed Identity token error: {e}")
|
|
return None
|
|
|
|
async def _get_certificate_token(self, scope: str) -> str | None:
|
|
try:
|
|
from cryptography.hazmat.primitives import serialization
|
|
from cryptography.hazmat.backends import default_backend
|
|
import jwt as pyjwt
|
|
|
|
cert_path = Path(self.certificate_path)
|
|
if not cert_path.exists():
|
|
logger.error(f"Certificate not found: {self.certificate_path}")
|
|
return None
|
|
|
|
with open(cert_path, "rb") as f:
|
|
private_key = serialization.load_pem_private_key(f.read(), password=None, backend=default_backend())
|
|
|
|
token_url = TOKEN_URL_TEMPLATE.format(tenant=self.tenant_id)
|
|
now = int(time.time())
|
|
assertion = pyjwt.encode(
|
|
{
|
|
"aud": token_url,
|
|
"iss": self.client_id,
|
|
"sub": self.client_id,
|
|
"jti": os.urandom(16).hex(),
|
|
"nbf": now,
|
|
"exp": now + 600,
|
|
},
|
|
private_key,
|
|
algorithm="RS256",
|
|
headers={"x5t": self.certificate_thumbprint} if self.certificate_thumbprint else {},
|
|
)
|
|
|
|
data = {
|
|
"client_id": self.client_id,
|
|
"client_assertion": assertion,
|
|
"client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
|
|
"scope": scope,
|
|
"grant_type": "client_credentials",
|
|
}
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.post(token_url, data=data) as resp:
|
|
if resp.status == 200:
|
|
result = await resp.json()
|
|
return result.get("access_token")
|
|
body = await resp.text()
|
|
logger.warning(f"Certificate token failed: HTTP {resp.status} - {body[:200]}")
|
|
except ImportError:
|
|
logger.error("cryptography and PyJWT required for certificate auth")
|
|
except Exception as e:
|
|
logger.error(f"Certificate token error: {e}")
|
|
return None
|
|
|
|
|
|
class DelegatedAuthStore:
|
|
_KEY_ENV = "MSTEAMS_DELEGATED_STORE_KEY"
|
|
|
|
def __init__(self, storage_dir: str | None = None):
|
|
self._storage_dir = Path(storage_dir or str(Path.home() / ".yuxi" / "msteams"))
|
|
self._storage_dir.mkdir(parents=True, exist_ok=True)
|
|
self._tokens: dict[str, dict[str, Any]] = {}
|
|
self._fernet = self._init_fernet()
|
|
self._load()
|
|
|
|
@property
|
|
def file_path(self) -> Path:
|
|
return self._storage_dir / DELEGATED_STORE_FILENAME
|
|
|
|
def _init_fernet(self):
|
|
from cryptography.fernet import Fernet
|
|
|
|
key_b64 = os.environ.get(self._KEY_ENV)
|
|
if key_b64:
|
|
return Fernet(key_b64.encode())
|
|
|
|
key_path = self._storage_dir / ".store_key"
|
|
if key_path.exists():
|
|
return Fernet(key_path.read_bytes())
|
|
|
|
key = Fernet.generate_key()
|
|
key_path.write_bytes(key)
|
|
return Fernet(key)
|
|
|
|
def _load(self) -> None:
|
|
if not self.file_path.exists():
|
|
return
|
|
try:
|
|
data = self.file_path.read_bytes()
|
|
try:
|
|
plain = self._fernet.decrypt(data).decode("utf-8")
|
|
except Exception:
|
|
plain = data.decode("utf-8")
|
|
self._save()
|
|
self._tokens = json.loads(plain).get("tokens", {})
|
|
except (json.JSONDecodeError, OSError) as e:
|
|
logger.error(f"MSTeams delegated store: failed to load: {e}")
|
|
|
|
def _save(self) -> None:
|
|
try:
|
|
plain = json.dumps({"tokens": self._tokens}, ensure_ascii=False, indent=2)
|
|
encrypted = self._fernet.encrypt(plain.encode("utf-8"))
|
|
self.file_path.write_bytes(encrypted)
|
|
except OSError as e:
|
|
logger.error(f"MSTeams delegated store: failed to save: {e}")
|
|
|
|
def store_token(self, user_id: str, access_token: str, refresh_token: str, expires_in: int = 3600) -> None:
|
|
self._tokens[user_id] = {
|
|
"access_token": access_token,
|
|
"refresh_token": refresh_token,
|
|
"expires_at": time.time() + expires_in,
|
|
}
|
|
self._save()
|
|
|
|
def get_token(self, user_id: str) -> dict[str, Any] | None:
|
|
entry = self._tokens.get(user_id)
|
|
if not entry:
|
|
return None
|
|
return entry
|
|
|
|
async def refresh_token(
|
|
self,
|
|
user_id: str,
|
|
client_id: str,
|
|
client_secret: str,
|
|
) -> str | None:
|
|
entry = self._tokens.get(user_id)
|
|
if not entry or not entry.get("refresh_token"):
|
|
return None
|
|
|
|
token_url = TOKEN_URL_TEMPLATE.format(tenant="common")
|
|
data = {
|
|
"client_id": client_id,
|
|
"client_secret": client_secret,
|
|
"refresh_token": entry["refresh_token"],
|
|
"grant_type": "refresh_token",
|
|
"scope": GRAPH_SCOPE,
|
|
}
|
|
try:
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.post(token_url, data=data) as resp:
|
|
if resp.status == 200:
|
|
result = await resp.json()
|
|
new_token = result.get("access_token", "")
|
|
new_refresh = result.get("refresh_token", "")
|
|
if new_token:
|
|
self.store_token(
|
|
user_id,
|
|
new_token,
|
|
new_refresh or entry["refresh_token"],
|
|
result.get("expires_in", 3600),
|
|
)
|
|
return new_token
|
|
logger.warning(f"Delegated token refresh failed: HTTP {resp.status}")
|
|
except Exception as e:
|
|
logger.error(f"Delegated token refresh error: {e}")
|
|
return None
|
|
|
|
def is_expired(self, user_id: str) -> bool:
|
|
entry = self._tokens.get(user_id)
|
|
if not entry:
|
|
return True
|
|
return time.time() >= entry.get("expires_at", 0) - DELEGATED_REFRESH_BUFFER_S
|